Update app.py

app.py

@@ -160,71 +160,48 @@
 # # Load any additional models if needed
 # # gr.load("models/Bhaskar2611/Capstone").launch()
 
-import torch
-import transformers
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Load tokenizer and model from Hugging Face (your fine-tuned model)
+# Load your Hugging Face token from environment variables
+hf_token = os.environ.get("HF_TOKEN")
+
 model_id = "Bhaskar2611/Capstone"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    use_auth_token=True
+# Load tokenizer and model with token parameter (no deprecated args)
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token)
+
+# Define a skin assistant prompt to set the context for the model
+SKIN_ASSISTANT_PROMPT = (
+    "You are a helpful assistant specialized in skin diseases and dermatology. "
+    "Provide accurate, concise, and helpful advice about skin conditions, symptoms, "
+    "treatments, and care. Always respond in a clear and empathetic way.\n\n"
 )
 
-streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-# Function to format prompt with Zephyr-style and system prompt
-def format_prompt(message, history):
-    system_prompt = (
-        "You are a helpful medical assistant specialized in skin diseases. "
-        "Always provide accurate, responsible, and actionable information. "
-        "If needed, recommend seeing a dermatologist."
-    )
-    prompt = f"<|user|>\n{system_prompt}\n<|assistant|>\nSure, I'm here to help you.</s>\n"
-    for user_msg, bot_msg in history:
-        prompt += f"<|user|>\n{user_msg}\n<|assistant|>\n{bot_msg}</s>\n"
-    prompt += f"<|user|>\n{message}\n<|assistant|>\n"
-    return prompt
-
-# Function to generate model response
-def respond(message, history):
-    prompt = format_prompt(message, history)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    output = model.generate(
-        **inputs,
-        max_new_tokens=2048,
-        pad_token_id=tokenizer.eos_token_id,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        streamer=streamer
-    )
-    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
-    # Extract the final assistant message only
-    final_response = decoded_output.split("<|assistant|>")[-1].strip().split("</s>")[0].strip()
-    return final_response
-
-# Create Gradio interface
-chat = gr.ChatInterface(
-    fn=respond,
-    title="Skin & Hair Disease Assistant",
-    chatbot=gr.Chatbot(height=400),
-    textbox=gr.Textbox(placeholder="Describe your symptoms or ask a question...", container=False, scale=7),
-    description="Ask about skin conditions, hair issues, or treatment guidance. Powered by a custom fine-tuned Zephyr model.",
-    theme="soft",
-    examples=["I have red patches on my skin.", "What causes hair loss?", "Is dandruff a fungal infection?"],
-    cache_examples=False,
-    retry_btn="🔁 Retry",
-    undo_btn="↩️ Undo",
-    clear_btn="🧹 Clear"
+def generate_response(user_input):
+    # Combine the assistant prompt + user input
+    prompt = SKIN_ASSISTANT_PROMPT + user_input
+    
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=2048, do_sample=True, temperature=0.7)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Remove the assistant prompt part from the response (if it appears)
+    if response.startswith(SKIN_ASSISTANT_PROMPT):
+        response = response[len(SKIN_ASSISTANT_PROMPT):].strip()
+    return response
+
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(lines=3, placeholder="Ask about skin diseases..."),
+    outputs="text",
+    title="Skin Disease Assistant",
+    description="Ask any questions related to skin diseases and get expert-like responses."
 )
 
-# Launch app
 if __name__ == "__main__":
-    chat.launch()
+    iface.launch()