Amrender commited on
Commit
f7d54db
·
verified ·
1 Parent(s): 2f9524f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +122 -0
app.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import re
4
+ import gradio as gr
5
+ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
6
+ from peft import PeftModel
7
+ from threading import Thread
8
+
9
+ # ==========================================
10
+ # 1. SETUP & AUTHENTICATION
11
+ # ==========================================
12
+ # Mistral requires an HF token to download the base model.
13
+ # You must add your token in the HF Spaces "Settings" -> "Variables and secrets" as 'HF_TOKEN'
14
+ hf_token = os.environ.get("HF_TOKEN")
15
+
16
+ BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
17
+ ADAPTER_REPO = "your-username/Medical-Mistral-7B-LoRA" # <--- CHANGE THIS TO YOUR REPO
18
+
19
+ print("Booting up Clinical AI Server...")
20
+
21
+ tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=hf_token)
22
+ tokenizer.pad_token = tokenizer.eos_token
23
+
24
+ # Load Base Model in 4-bit precision
25
+ bnb_config = BitsAndBytesConfig(
26
+ load_in_4bit=True,
27
+ bnb_4bit_quant_type="nf4",
28
+ bnb_4bit_compute_dtype=torch.float16
29
+ )
30
+
31
+ base_model = AutoModelForCausalLM.from_pretrained(
32
+ BASE_MODEL,
33
+ quantization_config=bnb_config,
34
+ device_map="auto",
35
+ token=hf_token
36
+ )
37
+
38
+ # Merge Base Model with your Trained LoRA Adapter
39
+ print("Downloading and attaching LoRA Adapter from Hub...")
40
+ model = PeftModel.from_pretrained(base_model, ADAPTER_REPO, token=hf_token)
41
+ model.eval()
42
+ print("✅ Model Ready!")
43
+
44
+ # ==========================================
45
+ # 2. FORMATTING FUNCTION
46
+ # ==========================================
47
+ def format_and_clean(text):
48
+ end_triggers = ["regards", "hope this", "hope i", "let me know", "dr.", "thanks for"]
49
+ lower_text = text.lower()
50
+ cutoff = len(text)
51
+
52
+ for trigger in end_triggers:
53
+ idx = lower_text.find(trigger)
54
+ if idx != -1 and idx < cutoff:
55
+ cutoff = idx
56
+
57
+ clean = text[:cutoff].strip().replace('*', '')
58
+
59
+ sentences = clean.split(". ")
60
+ capitalized = [s.capitalize() for s in sentences if s]
61
+ clean = ". ".join(capitalized)
62
+
63
+ clean = re.sub(r'(?i)assessment\s*:', '**Assessment:** ', clean)
64
+ clean = re.sub(r'(?i)analysis\s*:', '\n\n**Analysis:** ', clean)
65
+ clean = re.sub(r'(?i)recommended action\s*:', '\n\n**Recommended Action:** ', clean)
66
+
67
+ return clean
68
+
69
+ # ==========================================
70
+ # 3. CHAT ENGINE & UI
71
+ # ==========================================
72
+ def clinical_chat(message, history):
73
+ prompt = f"""<s>[INST] You are a highly intelligent clinical AI. The user says: "{message}"
74
+
75
+ If the user provides symptoms, diagnose them. If the user asks about a known condition, provide treatment advice.
76
+
77
+ You MUST format your response exactly like this:
78
+ **Assessment:** [Name of the condition]
79
+ **Analysis:** [Brief explanation]
80
+ **Recommended Action:** [Medical advice]
81
+
82
+ Do not include greetings, sign-offs, or links. [/INST]"""
83
+
84
+ inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
85
+ streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=10.0)
86
+
87
+ generation_kwargs = dict(
88
+ **inputs,
89
+ streamer=streamer,
90
+ max_new_tokens=300,
91
+ temperature=0.2,
92
+ repetition_penalty=1.15,
93
+ do_sample=True,
94
+ pad_token_id=tokenizer.eos_token_id
95
+ )
96
+
97
+ thread = Thread(target=model.generate, kwargs=generation_kwargs)
98
+ thread.start()
99
+
100
+ partial_text = ""
101
+ for new_token in streamer:
102
+ partial_text += new_token
103
+ yield format_and_clean(partial_text)
104
+
105
+ # Launch UI inside Docker container
106
+ theme = gr.themes.Soft(primary_hue="blue", neutral_hue="slate")
107
+
108
+ demo = gr.ChatInterface(
109
+ fn=clinical_chat,
110
+ title="⚕️ Clinical AI Diagnostic Assistant",
111
+ description="**Enter your symptoms or medical queries below for a professional analysis.**",
112
+ theme=theme,
113
+ examples=[
114
+ "I am having severe headache, body pain and strain in my neck.",
115
+ "What medicine should I take if I have high cholesterol?",
116
+ "I have been sneezing and coughing for 3 days."
117
+ ]
118
+ )
119
+
120
+ if __name__ == "__main__":
121
+ # 0.0.0.0 and port 7860 are strictly required for Docker on Hugging Face Spaces
122
+ demo.launch(server_name="0.0.0.0", server_port=7860)