Update handler.py

handler.py

@@ -1,6 +1,6 @@
 from ctransformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import pipeline
-
+import json
 
 class EndpointHandler:
     def __init__(self, model_dir):
@@ -13,14 +13,13 @@ class EndpointHandler:
         self.model = AutoModelForCausalLM.from_pretrained(
             f"{self.model_dir}/comic_mistral-v5.2.q5_0.gguf",
             model_type="mistral",
-            lib="avx2",  # or "avx" or "basic" depending on your CPU
-            gpu_layers=0,  # Set to a number > 0 if you want to use GPU
+            lib="avx2",
+            gpu_layers=0,
             hf=True
         )
         self.tokenizer = AutoTokenizer.from_pretrained(self.model)
 
     def preprocess(self, data):
-        # You can add any preprocessing steps here if needed
         return data
 
     def __call__(self, data):
@@ -30,20 +29,23 @@ class EndpointHandler:
         inputs = self.preprocess(data)
         prompt = inputs["inputs"]
 
-        # Generate text using the model directly for streaming
+        # 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):
-            yield self.postprocess({"generated_text": text})
+            generated_text += text
+
+        # Return a JSON-serializable response
+        response = {"generated_text": generated_text}
+        return json.dumps(response)
 
     def postprocess(self, data):
-        # You can add any postprocessing steps here if needed
         return data
 
-
 def get_handler(model_dir):
     handler = EndpointHandler(model_dir)
     handler.load_model()