sdfdsfads23333

app.py

@@ -12,7 +12,7 @@ from transformers import AutoModel, AutoProcessor, AutoTokenizer, AutoModelForCa
 OCR_MODEL_ID = "ibm-granite/granite-docling-258M"
 
 # Stage 2: LLM 모델 (텍스트에서 약 이름 추출)
-LLM_MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
+LLM_MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
 
 
 def _load_ocr_model():
@@ -50,7 +50,7 @@ print("🔄 Loading Granite Docling OCR model...")
 OCR_MODEL, OCR_PROCESSOR = _load_ocr_model()
 print("✅ OCR model loaded!")
 
-print("🔄 Loading Llama-3.1-8B-Instruct...")
+print("🔄 Loading Qwen2.5-7B-Instruct...")
 LLM_MODEL, LLM_TOKENIZER = _load_llm_model()
 print("✅ LLM model loaded!")
 
@@ -88,21 +88,24 @@ def extract_text_from_image(image: Image.Image) -> str:
 
 
 def extract_medications_from_text(text: str) -> List[str]:
-    """Stage 2: Llama 3.1로 텍스트에서 약 이름만 추출"""
+    """Stage 2: Qwen2.5로 텍스트에서 약 이름만 추출"""
     try:
-        prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
-
-You are a medical text analyzer. Extract only medication names from the given text and return them as a JSON array.
-Return ONLY valid JSON format: {{"medications": ["name1", "name2"]}}
-<|eot_id|><|start_header_id|>user<|end_header_id|>
-
-Extract all medication names from this text:
-
-{text}
-
-Return only the JSON array of medication names.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
-
-"""
+        messages = [
+            {
+                "role": "system",
+                "content": "You are a medical text analyzer. Extract only medication names from the given text and return them as a JSON array. Return ONLY valid JSON format."
+            },
+            {
+                "role": "user",
+                "content": f"Extract all medication names from this text:\n\n{text}\n\nReturn format: {{\"medications\": [\"name1\", \"name2\"]}}"
+            }
+        ]
+
+        prompt = LLM_TOKENIZER.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
 
         inputs = LLM_TOKENIZER(prompt, return_tensors="pt").to(LLM_MODEL.device)
 
@@ -118,9 +121,10 @@ Return only the JSON array of medication names.<|eot_id|><|start_header_id|>assi
 
         response = LLM_TOKENIZER.decode(outputs[0], skip_special_tokens=True)
 
-        # Extract assistant response
-        if "<|start_header_id|>assistant<|end_header_id|>" in response:
-            response = response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
+        # Extract assistant response (Qwen format)
+        if "<|im_start|>assistant" in response:
+            response = response.split("<|im_start|>assistant")[-1]
+            response = response.replace("<|im_end|>", "").strip()
 
         # Parse JSON
         json_match = re.search(r'\{.*?\}', response, re.DOTALL)
@@ -305,7 +309,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
 
     **ℹ️ 2단계 파이프라인**
     - **Stage 1**: Granite Docling (OCR) - 이미지에서 모든 텍스트 추출
-    - **Stage 2**: Llama 3.1 8B (LLM) - 추출된 텍스트에서 약 이름만 식별
+    - **Stage 2**: Qwen2.5 7B (LLM) - 추출된 텍스트에서 약 이름만 식별
 
     실제 복약은 의사·약사의 지시를 따르세요.
     """)