Update app.py

app.py

@@ -160,61 +160,115 @@
 # # Load any additional models if needed
 # # gr.load("models/Bhaskar2611/Capstone").launch()
 
-import os
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-import gradio as gr
+# import os
+# from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+# import gradio as gr
 
-# Load your Hugging Face token (if needed for private models or API limit increases)
-hf_token = os.environ.get("HF_TOKEN")
+# # Load your Hugging Face token (if needed for private models or API limit increases)
+# hf_token = os.environ.get("HF_TOKEN")
 
-# Model ID for Mistral 7B Instruct
-model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+# # Model ID for Mistral 7B Instruct
+# model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+# # Load tokenizer
+# tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
 
-# BitsAndBytesConfig for 4-bit quantization to reduce memory usage
-bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+# # BitsAndBytesConfig for 4-bit quantization to reduce memory usage
+# bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+
+# # Load model with quantization and device mapping
+# model = AutoModelForCausalLM.from_pretrained(
+#     model_id,
+#     quantization_config=bnb_config,
+#     device_map="auto",
+#     token=hf_token
+# )
 
-# Load model with quantization and device mapping
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    quantization_config=bnb_config,
-    device_map="auto",
-    token=hf_token
+# # Skin assistant system prompt
+# 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"
+# )
+
+# def generate_response(user_input):
+#     prompt = SKIN_ASSISTANT_PROMPT + user_input
+#     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+#     outputs = model.generate(
+#         **inputs,
+#         max_new_tokens=1024,
+#         do_sample=True,
+#         temperature=0.7,
+#         top_p=0.95,
+#         repetition_penalty=1.1
+#     )
+#     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+#     return response.replace(SKIN_ASSISTANT_PROMPT, "").strip()
+
+# # 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."
+# )
+
+# if __name__ == "__main__":
+#     iface.launch()
+
+import os
+import gradio as gr
+from huggingface_hub import InferenceClient
+from dotenv import load_dotenv
+
+# Load API token from .env or environment
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")  # or directly use your token here
+
+# Initialize the Hugging Face inference client
+client = InferenceClient(
+    model="mistralai/Mistral-7B-Instruct-v0.3",
+    token=HF_TOKEN
 )
 
-# Skin assistant system prompt
+# Skin assistant prompt
 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"
+    "You are an AI Dermatologist chatbot designed to assist users with skin by only providing text "
+    "and if user information is not provided related to skin then ask what they want to know related to skin."
 )
 
-def generate_response(user_input):
-    prompt = SKIN_ASSISTANT_PROMPT + user_input
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=1024,
-        do_sample=True,
+def respond(message, history):
+    messages = [{"role": "system", "content": SKIN_ASSISTANT_PROMPT}]
+    for user_msg, bot_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    for chunk in client.chat.completions.create(
+        model="mistralai/Mistral-7B-Instruct-v0.3",
+        messages=messages,
+        max_tokens=1024,
         temperature=0.7,
         top_p=0.95,
-        repetition_penalty=1.1
-    )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response.replace(SKIN_ASSISTANT_PROMPT, "").strip()
-
-# 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."
+        stream=True,
+    ):
+        token = chunk.choices[0].delta.get("content", "")
+        response += token
+        yield response
+
+# Launch Gradio interface
+demo = gr.ChatInterface(
+    fn=respond,
+    title="skin-bot",
+    theme="default"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()
+