Update app.py

app.py

@@ -1,7 +1,80 @@
-from transformers import pipeline
+# 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
+            )
 
-pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
+    # Define the action for the button click
+    generate_button.click(
+        fn=generate_text,
+        inputs=[user_input, prompt_template],
+        outputs=output_text
+    )
 
-demo = gr.Interface.from_pipeline(pipe)
 demo.launch()