Update llm.py

llm.py

@@ -1,50 +1,56 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from functools import partial
 import asyncio
+from functools import partial
+
 model_name = "microsoft/phi-3-mini-128k-instruct"
 
-# 8-bit quantization config
-quant_config = BitsAndBytesConfig(
-    load_in_8bit=True
-)
+# Quantization
+quant_config = BitsAndBytesConfig(load_in_8bit=True)
 
-# tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_name,local_files_only=True)
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    local_files_only=True
+)
 
-# model
+# Load model
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    quantization_config=quant_config,
-    local_files_only=True,
-    trust_remote_code=True
+    model_name,
+    device_map="auto",
+    quantization_config=quant_config,
+    local_files_only=True,
+    trust_remote_code=True
 )
-# Mock LLM function (replace with your actual LLM API call)
+
+# ✅ SYNC function (important)
+def execute_llm(model, tokenizer, prompt: str):
+    messages = [
+        {"role": "user", "content": prompt}
+    ]
+
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=300,   # keep safe
+            do_sample=True,
+            temperature=0.7
+        )
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+# ✅ ASYNC wrapper
 async def call_llm(prompt: str):
-    # Simulate LLM call in executor
-    llm_function_with_args=partial(execute_llm(model,tokenizer,prompt))
-    loop = asyncio.get_event_loop()
-    result = await loop.run_in_executor(None, llm_function_with_args)
-    return result
-
-async def execute_llm(model,tokenizer,prompt:str):
-    prompt="what is json give an example "
-    data=""
-    full_prompt=prompt+" "+data
-    messages = [
-        {"role": "user", "content":full_prompt }
-    ]
-
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        return_tensors="pt"
-        ).to(model.device)
-        
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens=500
-        )
-
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    loop = asyncio.get_event_loop()
+
+    func = partial(execute_llm, model, tokenizer, prompt)
+
+    result = await loop.run_in_executor(None, func)
+    return result