Update handler.py

handler.py

@@ -59,15 +59,15 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 #     return response
 
 class EndpointHandler:
-    def __init__(self):
+    def __init__(self,model_path:str):
         # Load processor and model
         self.PROCESSOR = AutoProcessor.from_pretrained(
-            "marutitecblic/HtmlTocode",
+            model_path,
             trust_remote_code=True,
             # token=API_TOKEN,
         )
         self.MODEL = AutoModelForCausalLM.from_pretrained(
-            "marutitecblic/HtmlTocode",
+            model_path,
             # token=API_TOKEN,
             trust_remote_code=True,
             torch_dtype=torch.bfloat16,
@@ -99,7 +99,7 @@ class EndpointHandler:
         # inputs = preprocess(model_inputs)
         generated_ids = self.MODEL.generate(**inputs, bad_words_ids=self.BAD_WORDS_IDS, max_length=4096)
         generated_text = self.PROCESSOR.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        return {"generated_text": generated_text}
+        return {"text": generated_text}
         # return {"text":prediction[0]}
     
     # @classmethod