Update handler.py

handler.py

@@ -1,52 +1,42 @@
-from ctransformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import pipeline
 import json
-
-class EndpointHandler:
-    def __init__(self, model_dir):
-        self.model_dir = model_dir
-        self.model = None
-        self.tokenizer = None
-        self.pipe = None
-
-    def load_model(self):
-        self.model = AutoModelForCausalLM.from_pretrained(
-            f"{self.model_dir}/comic_mistral-v5.2.q5_0.gguf",
-            model_type="mistral",
-            lib="avx2",
-            gpu_layers=0,
-            hf=True
-        )
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model)
-
-    def preprocess(self, data):
-        return data
-
-    def __call__(self, data):
-        if self.model is None or self.tokenizer is None:
-            self.load_model()
-
-        inputs = self.preprocess(data)
-        prompt = inputs["inputs"]
-
-        # Generate text using the model
-        generated_text = ""
-        for text in self.model(prompt,
-                               max_new_tokens=256,
-                               temperature=0.8,
-                               repetition_penalty=1.1,
-                               do_sample=True,
-                               stream=True):
-            generated_text += text
-
-        # Return a JSON-serializable response
-        response = {"generated_text": generated_text}
-        return json.dumps(response)
-
-    def postprocess(self, data):
-        return data
-
-def get_handler(model_dir):
-    handler = EndpointHandler(model_dir)
-    handler.load_model()
-    return handler
+import os
+from typing import Dict, List, Any
+from llama_cpp import Llama
+import gemma_tools as gem
+
+MAX_TOKENS = 512
+
+class EndpointHandler():
+    def __init__(self, data):
+        # Update the model path and filename with your ComicBot model
+        self.model = Llama.from_pretrained("njwright92/ComicBot_v.2-gguf", filename="ComicBot_v.2-q4_k_m.gguf", n_ctx=8192)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        args = gem.get_args_or_none(data)
+        fmat = "<startofturn>system\n{system_prompt} <endofturn>\n<startofturn>user\n{prompt} <endofturn>\n<startofturn>model"
+        print(args, fmat)
+        if not args[0]:
+            return {
+                "status": args["status"],
+                "message": args["description"]
+            }
+        try:
+            fmat = fmat.format(system_prompt=args["system_prompt"], prompt=args["inputs"])
+        except Exception as e:
+            return json.dumps({
+                "status": "error",
+                "reason": "invalid format"
+            })
+        
+        max_length = data.pop("max_length", 512)
+        try:
+            max_length = int(max_length)
+        except Exception as e:
+            return json.dumps({
+                "status": "error",
+                "reason": "max_length was passed as something that was absolutely not a plain old int"
+            })
+        
+        res = self.model(fmat, temperature=args["temperature"], top_p=args["top_p"], top_k=args["top_k"], max_tokens=max_length)
+
+        return res