YanBoChen commited on
Commit
7282410
·
1 Parent(s): c414f60

WIP(llm_clients): add MeditronClient for medical query processing with Hugging Face integration

Browse files
Files changed (3) hide show
  1. .env.example +2 -0
  2. requirements.txt +1 -0
  3. src/llm_clients.py +165 -0
.env.example ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ # .env.example document
2
+ HF_TOKEN=your_huggingface_token_here
requirements.txt CHANGED
@@ -56,6 +56,7 @@ pydub==0.25.1
56
  Pygments==2.19.2
57
  pyparsing==3.2.3
58
  python-dateutil==2.9.0.post0
 
59
  python-multipart==0.0.20
60
  pytz==2025.2
61
  PyYAML==6.0.2
 
56
  Pygments==2.19.2
57
  pyparsing==3.2.3
58
  python-dateutil==2.9.0.post0
59
+ python-dotenv==1.1.1
60
  python-multipart==0.0.20
61
  pytz==2025.2
62
  PyYAML==6.0.2
src/llm_clients.py ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ OnCall.ai LLM Clients Module
3
+
4
+ Provides specialized LLM clients for medical query processing.
5
+
6
+ Author: OnCall.ai Team
7
+ Date: 2025-07-29
8
+ """
9
+
10
+ import logging
11
+ import os
12
+ from typing import Dict, Optional
13
+ from huggingface_hub import InferenceClient
14
+ from dotenv import load_dotenv
15
+
16
+ # Load environment variables from .env file
17
+ load_dotenv()
18
+
19
+ class MeditronClient:
20
+ def __init__(
21
+ self,
22
+ model: str = "TheBloke/meditron-7B-GPTQ",
23
+ timeout: float = 30.0
24
+ ):
25
+ """
26
+ Initialize Meditron API client for medical query processing.
27
+
28
+ Args:
29
+ model: Hugging Face model name
30
+ timeout: API call timeout duration (not used in InferenceClient)
31
+
32
+ Warning: This model should not be used for professional medical advice.
33
+ """
34
+ # Get HF token from environment variable
35
+ hf_token = os.getenv('HF_TOKEN')
36
+ if not hf_token:
37
+ raise ValueError(
38
+ "HF_TOKEN not found in environment variables. "
39
+ "Please set HF_TOKEN in your .env file or environment."
40
+ )
41
+
42
+ self.client = InferenceClient(model=model, token=hf_token)
43
+ self.logger = logging.getLogger(__name__)
44
+ self.timeout = timeout
45
+ self.logger.warning(
46
+ "Meditron Model: Research tool only. "
47
+ "Not for professional medical diagnosis."
48
+ )
49
+ self.logger.info("Meditron client initialized with HF token")
50
+
51
+ def analyze_medical_query(
52
+ self,
53
+ query: str,
54
+ max_tokens: int = 100,
55
+ timeout: Optional[float] = None
56
+ ) -> Dict[str, str]:
57
+ """
58
+ Analyze medical query and extract condition.
59
+
60
+ Args:
61
+ query: Medical query text
62
+ max_tokens: Maximum tokens to generate
63
+ timeout: Specific API call timeout (not used in InferenceClient)
64
+
65
+ Returns:
66
+ Extracted medical condition information
67
+ """
68
+ try:
69
+ # ChatML style prompt for Meditron
70
+ prompt = f"""<|im_start|>system
71
+ You are a professional medical assistant trained to extract medical conditions.
72
+ Provide only the most representative condition name.
73
+ DO NOT provide medical advice.
74
+ <|im_end|>
75
+ <|im_start|>user
76
+ {query}
77
+ <|im_end|>
78
+ <|im_start|>assistant
79
+ """
80
+
81
+ self.logger.info(f"Calling Meditron API with query: {query}")
82
+
83
+ # Remove timeout parameter as InferenceClient doesn't support it
84
+ response = self.client.text_generation(
85
+ prompt,
86
+ max_new_tokens=max_tokens,
87
+ temperature=0.7,
88
+ top_k=50
89
+ )
90
+
91
+ self.logger.info(f"Received response: {response}")
92
+
93
+ # Extract condition from response
94
+ extracted_condition = self._extract_condition(response)
95
+
96
+ return {
97
+ 'extracted_condition': extracted_condition,
98
+ 'confidence': 0.8,
99
+ 'raw_response': response
100
+ }
101
+
102
+ except Exception as e:
103
+ self.logger.error(f"Meditron API query error: {str(e)}")
104
+ self.logger.error(f"Error type: {type(e).__name__}")
105
+ return {
106
+ 'extracted_condition': '',
107
+ 'confidence': 0,
108
+ 'error': str(e)
109
+ }
110
+
111
+ def _extract_condition(self, response: str) -> str:
112
+ """
113
+ Extract medical condition from model response.
114
+
115
+ Args:
116
+ response: Full model-generated text
117
+
118
+ Returns:
119
+ Extracted medical condition
120
+ """
121
+ from medical_conditions import CONDITION_KEYWORD_MAPPING
122
+
123
+ # Remove prompt parts, keep only generated content
124
+ generated_text = response.split('<|im_start|>assistant\n')[-1].strip()
125
+
126
+ # Search in known medical conditions
127
+ for condition in CONDITION_KEYWORD_MAPPING.keys():
128
+ if condition.lower() in generated_text.lower():
129
+ return condition
130
+
131
+ return generated_text.split('\n')[0].strip()
132
+
133
+ def main():
134
+ """
135
+ Test Meditron client functionality
136
+ """
137
+ try:
138
+ client = MeditronClient()
139
+ test_queries = [
140
+ "patient experiencing chest pain",
141
+ "sudden weakness on one side",
142
+ "severe headache with neurological symptoms"
143
+ ]
144
+
145
+ for query in test_queries:
146
+ print(f"\nTesting query: {query}")
147
+ result = client.analyze_medical_query(query)
148
+ print("Extracted Condition:", result['extracted_condition'])
149
+ print("Confidence:", result['confidence'])
150
+ if 'error' in result:
151
+ print("Error:", result['error'])
152
+ print("---")
153
+
154
+ except Exception as e:
155
+ print(f"Client initialization error: {str(e)}")
156
+ print("This might be due to:")
157
+ print("1. Missing Hugging Face API token")
158
+ print("2. Network connectivity issues")
159
+ print("3. Model access permissions")
160
+ print("\nTo fix:")
161
+ print("1. Set HF_TOKEN environment variable")
162
+ print("2. Or login with: huggingface-cli login")
163
+
164
+ if __name__ == "__main__":
165
+ main()