Upload handler.py

handler.py

@@ -16,6 +16,7 @@ class EndpointHandler():
         config = PeftConfig.from_pretrained(path)
         model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, trust_remote_code=True, device_map='auto')
         self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+        self.tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'system' %}{{ bos_token + '<<SYS>>\\n' + message['content'] + '\\n<</SYS>>\\n\\n' }}{% elif message['role'] == 'assistant' %}{{ ' '  + message['content'] + ' ' + eos_token }}{% endif %}{% endfor %}"
         # Load the Lora model
         self.model = PeftModel.from_pretrained(model, path)
 
@@ -44,6 +45,8 @@ class EndpointHandler():
 
         # Call the generate function
         output = generate(
+            tokenizer=tokenizer,
+            model=model,
             message=message,
             chat_history=chat_history,
             system_prompt=system_prompt,
@@ -63,6 +66,8 @@ class EndpointHandler():
         return {"generated_text": prediction}
 
 def generate(
+    tokenizer,
+    model,
     message: str,
     chat_history: list[tuple[str, str]],
     system_prompt: str = None,