app.py

import gradio as gr
from transformers import AutoTokenizer, AutoProcessor, VisionEncoderDecoderModel, TrOCRProcessor
from vllm import LLM, SamplingParams
from PIL import Image

# Load the language model and tokenizer from Hugging Face
model_name = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Initialize vLLM with CPU configuration
vllm_model = LLM(model=model_name, tensor_parallel_size=1, device="cpu")

 

def generate_response(prompt, max_tokens, temperature, top_p):
    # Define sampling parameters
    sampling_params = SamplingParams(
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )
    
    # Generate text using vLLM (input is the raw string `prompt`)
    output = vllm_model.generate(prompt, sampling_params)
    
    # Extract and decode the generated tokens
    generated_text = output[0].outputs[0].text
    return generated_text

 

# Gradio UI
with gr.Blocks() as demo:
    
        
        submit_button.click(
            generate_response,
            inputs=[prompt_input, max_tokens, temperature, top_p],
            outputs=output_text,
        )

    with gr.Tab("OCR"):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    label="Upload Image",
                    type="filepath",  # Corrected type
                    image_mode="RGB",
                )
                ocr_submit_button = gr.Button("Extract Text")
            
            with gr.Column():
                ocr_output = gr.Textbox(
                    label="Extracted Text",
                    lines=10,
                    interactive=False,
                )
        
        ocr_submit_button.click(
            ocr_image,
            inputs=[image_input],
            outputs=ocr_output,
        )
prompt =gr.Textbox()
max_tokens = gr.Textbox()
temperature = gr.Textbox()
top_p = gr.Textbox()
demo=gr.Interface(generate_response, inputs=[prompt, max_tokens,temperature, top_p], outputs="text")

# Launch the app
demo.launch()