Create app.py

app.py

@@ -1,15 +1,28 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-import torch
+import requests
+import json
 
-model_name = "MBZUAI/BiMediX2-8B-hf"
+API_URL = "http://localhost:8000/v1/chat/completions"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
+def chat_fn(message, history):
+    headers = {"Content-Type": "application/json"}
+    payload = {
+        "model": "MBZUAI/BiMediX2-8B-hf",
+        "messages": [{"role": "user", "content": message}],
+        "max_tokens": 512,
+        "temperature": 0.2
+    }
 
-def chat(text):
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=300)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    try:
+        response = requests.post(API_URL, headers=headers, json=payload)
+        data = response.json()
+        reply = data["choices"][0]["message"]["content"]
+        return reply
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-gr.Interface(chat, inputs="text", outputs="text").launch()
+gr.ChatInterface(
+    fn=chat_fn,
+    title="BiMediX2 Medical Chatbot",
+    description="Ask any medical question."
+).launch()