Update llm.py

llm.py

@@ -1,5 +1,5 @@
-import os
-hf_token=os.getenv("HUGGINGFACE_HUB_TOKEN")
+from huggingface_hub import login
+login(token=os.getenv("HUGGINGFACE_HUB_TOKEN"))
 import torch
 import asyncio
 from functools import partial
@@ -8,16 +8,17 @@ from transformers import AutoProcessor, AutoModelForImageTextToText #, BitsAndBy
 
 # Quantization config
 # quant_config = BitsAndBytesConfig(load_in_8bit=True)
-
+model_name="Qwen/Qwen3.5-9B-Base"
 # Load processor
-processor = AutoProcessor.from_pretrained("datalab-to/chandra-ocr-2")
+processor = AutoProcessor.from_pretrained(model_name)
 
 # Load model (auto device mapping)
 model = AutoModelForImageTextToText.from_pretrained(
-    "datalab-to/chandra-ocr-2",
+    model_name,
     # quantization_config=quant_config,
     device_map="auto",
-    token=hf_token
+    attn_implementation='flash_attention_2'
+    
 )
 
 print("CUDA available:", torch.cuda.is_available())
@@ -29,26 +30,21 @@ if torch.cuda.is_available():
 
     
 def execute_llm(model, processor, image, prompt: str):
-    print("execute llm")
-    print(prompt)
-    a=time.time()
-    # ✅ Use passed prompt (FIXED)
+
     if not prompt:
         prompt = """
-        Extract all text from the given image and return key value pair like json.
+        Extract all text from image.
+        Return ONLY valid JSON.
         """
 
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompt}
-            ]
-        }
-    ]
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": prompt}
+        ]
+    }]
 
-    
     with torch.inference_mode():
 
         inputs = processor.apply_chat_template(
@@ -57,31 +53,25 @@ def execute_llm(model, processor, image, prompt: str):
             tokenize=True,
             return_dict=True,
             return_tensors="pt"
-        ).to(model.device)
+        )
 
-        print("inputs ready")
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         outputs = model.generate(
             **inputs,
-            max_new_tokens=1000,   # 🔥 REDUCED (important)
-            do_sample=False       # deterministic output
+            max_new_tokens=200,
+            do_sample=False
         )
 
-        print("generated")
-
-        result = processor.decode(
+        return processor.decode(
             outputs[0][inputs["input_ids"].shape[-1]:],
             skip_special_tokens=True
         )
-    print('total time taken',time.time()-a)
-    print(result)
-    return result
-
 
 async def call_llm(image, prompt: str = ""):
     print("call llm")
 
-
-    result=execute_llm(model,processor,image,prompt)
+    loop = asyncio.get_event_loop()
+    result = await loop.run_in_executor(None, execute_llm, model, processor, image, prompt)
 
     return result