Update model_runner.py

model_runner.py

@@ -1,22 +1,40 @@
-import os
-from llama_cpp import Llama  # assuming llama-cpp-python is used to run GGUF
+import json
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
+from llama_cpp_agent.providers import LlamaCppPythonProvider
 
 class ModelRunner:
-    def __init__(self, model_path: str):
-        # Initialize the model once
-        self.llm = Llama(model_path=model_path)
+    def __init__(self, repo_id: str, filename: str):
+        model_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            local_dir="./models"
+        )
+        self.llm = Llama(
+            model_path=model_path,
+            flash_attn=True,
+            n_gpu_layers=81,
+            n_batch=1024,
+            n_ctx=8192,
+        )
+        self.provider = LlamaCppPythonProvider(self.llm)
+        self.agent = LlamaCppAgent(
+            self.provider,
+            system_prompt="You are an AI developer. Given a codegen output, return a JSON mapping of filenames to file contents.",
+            predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
+            debug_output=False
+        )
 
-    def interpret_code_description(self, prompt: str) -> str:
-        """
-        Given the pasted AI response (prompt), return the model’s parsing output
-        (e.g. JSON or a DSL).
-        """
-        # You may want to wrap the prompt with instructions
-        wrapped = (
-            "You are a code parser. The following is a description of multiple files with code. "
-            "Produce as output a JSON object mapping filenames to file contents. "
-            "If something is ambiguous, try your best.\n\n"
-            + prompt
+    def interpret_code_description(self, pasted_text: str) -> str:
+        user_prompt = (
+            "Here is an AI output that includes multiple code files. "
+            "Please convert it into a JSON object like {\"file1.py\": \"...code...\", \"file2.txt\": \"...\"}\n\n"
+            f"{pasted_text}"
         )
-        resp = self.llm(prompt=wrapped, max_tokens=2048)
-        return resp["choices"][0]["text"]
+
+        output = ""
+        for chunk in self.agent.get_chat_response(user_prompt, returns_streaming_generator=True):
+            output += chunk
+
+        return output