Fix gradio app to use actual TinyLlama model for skin disease assistance

Co-authored-by: Amp <[email protected]>
Amp-Thread-ID: https://ampcode.com/threads/T-31dcde95-a69d-4276-aff8-3096bca9d588

gradio-app.py

@@ -1,7 +1,57 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Load the model and tokenizer
+model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Add pad token if it doesn't exist
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+def skin_disease_assistant(user_input):
+    # Create a medical-focused prompt
+    prompt = f"""You are a medical AI assistant specializing in skin diseases. Provide helpful, accurate information about skin conditions based on the following query.
+
+Query: {user_input}
+
+Response:"""
+    
+    # Tokenize and generate response
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode the response
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Extract only the generated part (after "Response:")
+    response = full_response.split("Response:")[-1].strip()
+    
+    return response
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=skin_disease_assistant,
+    inputs=gr.Textbox(label="Ask about skin conditions", placeholder="e.g., What are the symptoms of eczema?"),
+    outputs=gr.Textbox(label="AI Response"),
+    title="🏥 Skin Disease AI Assistant",
+    description="An AI assistant to help with skin disease questions. For educational purposes only - consult medical professionals for actual diagnosis.",
+    examples=[
+        "What are the symptoms of psoriasis?",
+        "How to treat acne?",
+        "What causes eczema flare-ups?",
+        "Difference between melanoma and mole?"
+    ]
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()