# from transformers import pipeline
# import gradio as gr

# pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")

# demo = gr.Interface.from_pipeline(pipe)
# demo.launch()

import gradio as gr
from transformers import pipeline
import torch

# Define the model to use.
MODEL_NAME = "google/flan-t5-small"

# Set up the pipeline, specifying the task and the model.
# The pipeline handles tokenization and model inference.
# Using device="cuda" if a GPU is available, otherwise falls back to CPU.
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipeline("text2text-generation", model=MODEL_NAME, device=device)

# Define the function that constructs the prompt and calls the pipeline.
def generate_text(user_input, prompt_template):
    """
    Combines user input with a template and calls the Hugging Face transformers pipeline.
    """
    # Create the full prompt based on the template and user input.
    full_prompt = prompt_template.format(user_input=user_input)

    # Use the pipeline to generate text.
    try:
        # The pipeline returns a list of dictionaries; we extract the generated text.
        response = pipe(full_prompt, max_new_tokens=100)
        return response[0]['generated_text']
    except Exception as e:
        return f"Error: {e}"

# Define the Gradio interface.
with gr.Blocks() as demo:
    gr.Markdown("# Lightweight LLM Demo")
    gr.Markdown("Enter text and select a prompt to generate an AI response.")

    with gr.Row():
        with gr.Column(scale=1):
            # Textbox for user input
            user_input = gr.Textbox(
                label="Your Input Text",
                placeholder="Type here...",
                lines=5
            )

            # Dropdown to select a prompt template
            prompt_template = gr.Dropdown(
                label="Choose a Prompt Template",
                choices=[
                    "Summarize this: {user_input}",
                    "Answer the following question: {user_input}",
                    "Rewrite this text to be more formal: {user_input}"
                ],
                value="Summarize this: {user_input}"
            )

            # Button to trigger the generation
            generate_button = gr.Button("Generate")

        with gr.Column(scale=2):
            # Textbox to display the output
            output_text = gr.Textbox(
                label="Generated Output",
                lines=10
            )

    # Define the action for the button click
    generate_button.click(
        fn=generate_text,
        inputs=[user_input, prompt_template],
        outputs=output_text
    )

demo.launch()