Update app.py

app.py

@@ -1,59 +1,40 @@
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Specify the model name from HuggingFace Hub
-model_name = "Qwen/Qwen2.5-0.5B-Instruct"
-
-# Load the tokenizer and model.
-# Using device_map="auto" ensures that if a GPU is available (as on ZeroGPU), it is used.
+# Load tokenizer and model
+model_name = "Qwen/Qwen2.5-0.5B-Instruct"  # replace with actual model name or path
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 
-# Create a text-generation pipeline.
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+if torch.cuda.is_available():
+    # Load model in half precision and move to GPU for faster inference
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+    model.to("cuda")
+else:
+    # Load model on CPU (full precision by default)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def generate_recommendation(characteristics: str) -> str:
-    """
-    Generates personalized diet recommendations from the provided characteristics.
-    """
-    prompt = (
-        f"Based on the following characteristics, provide personalized diet recommendations:\n\n"
-        f"{characteristics}\n\n"
-        f"Diet Recommendations:"
-    )
-    try:
-        output = generator(prompt, max_length=512, num_return_sequences=1)
-        # Wrap the result in a paragraph tag for HTML rendering.
-        recommendation = f"<p>{output[0]['generated_text']}</p>"
-        return recommendation
-    except Exception as e:
-        return f"<p style='color: red;'>Error: {str(e)}</p>"
+# Define the text generation function
+def generate_text(prompt):
+    # Tokenize input and move to appropriate device
+    inputs = tokenizer(prompt, return_tensors="pt")
+    if torch.cuda.is_available():
+        inputs = {key: value.to("cuda") for key, value in inputs.items()}
+    # Generate text
+    output_ids = model.generate(**inputs, max_length=200, do_sample=True, temperature=0.7)
+    # Decode the generated tokens to text
+    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return generated_text
 
-# Build the Gradio Blocks interface.
+# Set up Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# AI Diet Recommendations")
-    gr.Markdown(
-        "Get personalized diet recommendations generated by our AI."
-    )
-    
-    with gr.Row():
-        with gr.Column():
-            characteristics = gr.Textbox(
-                lines=10,
-                placeholder="Enter your characteristics here...",
-                label="Your Characteristics"
-            )
-            submit = gr.Button("Submit")
-        with gr.Column():
-            result = gr.HTML(label="Diet Recommendations")
-    
-    # When the button is clicked, the app will show a progress spinner.
-    submit.click(
-        fn=generate_recommendation,
-        inputs=characteristics,
-        outputs=result,
-        show_progress=True
-    )
+    gr.Markdown("## Text Generation Demo")  # title or description (optional)
+    user_input = gr.Textbox(label="Enter your prompt:")
+    output_text = gr.Textbox(label="Generated output:")
+    generate_btn = gr.Button("Generate")
+    # When button is clicked, call generate_text with user_input and show in output_text
+    generate_btn.click(fn=generate_text, inputs=user_input, outputs=output_text)
 
+# Launch the app (if running as a standalone script)
 if __name__ == "__main__":
     demo.launch()