Update app.py

app.py

@@ -1,12 +1,15 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the TinyLlama model for text generation
-pipe = pipeline("text-generation", model="TinyLlama/TinyLlama_v1.1").to("cuda")
+# Load the TinyLlama model for text generation on GPU
+pipe = pipeline(
+    "text-generation",
+    model="TinyLlama/TinyLlama_v1.1",
+    device=0  # 0 for 'cuda:0', -1 for CPU
+)  # No .to("cuda") needed[4][6]
 
 # 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,
@@ -15,7 +18,6 @@ def generate_text(prompt, max_length=128, temperature=1.0, top_p=0.95):
         num_return_sequences=1,
         do_sample=True
     )
-    # The output is a list of dicts with 'generated_text'
     return result[0]['generated_text']
 
 # Create the Gradio interface