Update inference.py

inference.py

@@ -3,179 +3,83 @@ import json
 import re
 import logging
 import warnings
-from typing import Optional, Dict, Any
-
 from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
 
-
-# ---------------------------
-# Silence noisy logs
-# ---------------------------
+# Silence unnecessary logs
 warnings.filterwarnings("ignore", category=UserWarning)
 logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
-
 os.environ["GGML_PYTHON_VERBOSE"] = "0"
-os.environ["LLAMA_CPP_LIB_VERBOSE"] = "0"
-
 
-# ---------------------------
-# Core Class
-# ---------------------------
 class HealthFunctionLM:
-    def __init__(
-        self,
-        repo_id: Optional[str] = None,
-        filename: Optional[str] = None,
-        n_ctx: int = 2048,
-        n_threads: int = 4,
-        temperature: float = 0.1,
-    ):
-        """
-        Initialize model
-
-        Args:
-            repo_id: Hugging Face repo ID
-            filename: GGUF file name
-            n_ctx: context length
-            n_threads: CPU threads
-            temperature: sampling temperature
-        """
-
-        # Defaults (easy mode)
-        self.repo_id = repo_id or "ramgovindv/health_function_call_llama3.2_3b_gguf"
-        self.filename = filename or "Llama-3.2-3B-Instruct.Q4_K_M.gguf"
-        self.temperature = temperature
-
-        if not self.filename.endswith(".gguf"):
-            raise ValueError("Only GGUF models are supported")
-
-        self.llm = Llama.from_pretrained(
-            repo_id=self.repo_id,
-            filename=self.filename,
+    """
+    A specialized wrapper for Llama-3.2 GGUF models to perform 
+    health-related function calling.
+    """
+    def __init__(self, repo_id: str, filename: str, n_ctx: int = 2048):
+        # Download model automatically from Hugging Face
+        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        
+        self.llm = Llama(
+            model_path=model_path,
             n_ctx=n_ctx,
-            n_threads=n_threads,
-            chat_format=None,
-            verbose=False,
+            n_threads=os.cpu_count() or 4,
+            verbose=False
         )
 
-    # ---------------------------
-    # Prompt Builder
-    # ---------------------------
     def _build_prompt(self, query: str) -> str:
-        return f"""
-You are an API generator.
-
-Return ONLY valid JSON in this format:
-{{
-  "name": "function_name",
-  "parameters": {{}},
-  "reasoning": "optional short explanation"
-}}
-
-User query:
-{query}
-
-JSON:
-"""
-
-    # ---------------------------
-    # Model Call
-    # ---------------------------
-    def _generate(self, prompt: str) -> str:
-        response = self.llm.create_chat_completion(
-            messages=[{"role": "user", "content": prompt}],
-            temperature=self.temperature,
+        return (
+            "You are an API generator. Return ONLY a JSON object.\n"
+            "Format: {\"name\": \"function_name\", \"parameters\": {\"key\": \"value\"}}\n\n"
+            f"User query: {query}\n\n"
+            "JSON:"
         )
 
-        return response["choices"][0]["message"]["content"].strip()
-
-    # ---------------------------
-    # Parsing Utilities
-    # ---------------------------
-    def _safe_json_load(self, text: str) -> Optional[Dict[str, Any]]:
-        """
-        Try strict JSON parsing first.
-        Fallback: extract first JSON block.
-        """
-        try:
-            return json.loads(text)
-        except Exception:
-            pass
-
-        # fallback: extract JSON substring
+    def _extract_json(self, text: str):
+        """Extracts JSON even if the model wraps it in markdown blocks."""
+        # Remove <think> tags if present
+        text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
+        # Find JSON block
         match = re.search(r"\{.*\}", text, re.DOTALL)
         if match:
             try:
                 return json.loads(match.group(0))
-            except Exception:
+            except json.JSONDecodeError:
                 return None
-
         return None
 
-    # ---------------------------
-    # Public Query API
-    # ---------------------------
-    def query(self, user_query: str) -> Dict[str, Any]:
+    def query(self, user_query: str):
         prompt = self._build_prompt(user_query)
-        raw_output = self._generate(prompt)
-
-        parsed = self._safe_json_load(raw_output)
-
-        if parsed:
+        output = self.llm.create_chat_completion(
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.1
+        )
+        
+        message = output["choices"][0]["message"]
+        content = message.get("content", "").strip()
+        
+        # Try to parse function call
+        function_data = self._extract_json(content)
+        
+        if function_data:
             return {
                 "query": user_query,
                 "type": "function_call",
-                "data": {
-                    "name": parsed.get("name"),
-                    "parameters": parsed.get("parameters", {}),
-                    "reasoning": parsed.get("reasoning"),
-                },
-                "raw": raw_output,
+                "data": function_data
             }
 
-        # fallback (model messed up)
         return {
             "query": user_query,
             "type": "text",
-            "data": {
-                "content": raw_output,
-                "reasoning": None,
-            },
+            "data": {"content": content}
         }
 
-
-# ---------------------------
-# Simple Loader (User Entry)
-# ---------------------------
-def load_model(
-    repo_id: Optional[str] = None,
-    filename: Optional[str] = None,
-    **kwargs
-) -> HealthFunctionLM:
-    """
-    Easy entry point for users
-
-    Example:
-        model = load_model()
-        model = load_model(repo_id="other/repo", filename="model.gguf")
-    """
-    return HealthFunctionLM(
-        repo_id=repo_id,
-        filename=filename,
-        **kwargs
-    )
-
-
-# ---------------------------
-# Optional CLI usage
-# ---------------------------
+# --- Example Usage ---
 if __name__ == "__main__":
-    model = load_model()
-
-    while True:
-        q = input("\nEnter query (or 'exit'): ")
-        if q.lower() == "exit":
-            break
-
-        result = model.query(q)
-        print(json.dumps(result, indent=2))
+    model = HealthFunctionLM(
+        repo_id="ramgovindv/health_function_call_llama3.2_3b_gguf",
+        filename="Llama-3.2-3B-Instruct.Q4_K_M.gguf"
+    )
+    
+    res = model.query("I am feeling very dizzy for couple of days. what could be the reason")
+    print(json.dumps(res, indent=2))