Update app.py

app.py

@@ -1,43 +1,36 @@
 import gradio as gr
-import torch
 from transformers import pipeline
 
-# Model ID for Llama 3 8B instruct (replace with the exact model you want)
-MODEL_ID = "manycore-research/SpatialLM-Llama-1B"
+# Load the TinyLlama model for text generation
+pipe = pipeline("text-generation", model="TinyLlama/TinyLlama_v1.1")
 
-# Load the text-generation pipeline with device_map="auto" to use GPU if available
-generator = pipeline(
-    "text-generation",
-    model=MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-
-def generate_response(prompt, max_length=512, temperature=0.7):
-    # Format prompt for Llama 3 instruct style
-    formatted_prompt = f"<s>[INST] {prompt} [/INST]"
-    output = generator(
-        formatted_prompt,
+# Define the prediction function
+def generate_text(prompt, max_length=128, temperature=1.0, top_p=0.95):
+    # You can expose more parameters as needed
+    result = pipe(
+        prompt,
         max_length=max_length,
         temperature=temperature,
-        do_sample=True,
-        top_p=0.95,
+        top_p=top_p,
         num_return_sequences=1,
+        do_sample=True
     )
-    generated_text = output[0]["generated_text"]
-    # Extract the response after the [/INST] token
-    response = generated_text.split("[/INST]")[-1].strip()
-    return response
+    # The output is a list of dicts with 'generated_text'
+    return result[0]['generated_text']
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Chat with Llama 3 (8B Instruct)")
-    with gr.Row():
-        with gr.Column():
-            user_input = gr.Textbox(lines=3, placeholder="Type your message here...", label="Your Message")
-            submit_btn = gr.Button("Submit")
-        with gr.Column():
-            output = gr.Textbox(lines=10, label="Llama 3 Response")
-    submit_btn.click(fn=generate_response, inputs=user_input, outputs=output)
+# Create the Gradio interface
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(lines=4, label="Input Prompt"),
+        gr.Slider(32, 512, value=128, step=8, label="Max Length"),
+        gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature"),
+        gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top-p (nucleus sampling)")
+    ],
+    outputs=gr.Textbox(lines=8, label="Generated Text"),
+    title="TinyLlama Text Generation",
+    description="Enter a prompt and generate text using TinyLlama/TinyLlama_v1.1."
+)
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the app
+demo.launch()