Update handler.py

handler.py

@@ -1,22 +1,12 @@
-# Code to inference Hermes with HF Transformers
-# Requires pytorch, transformers, bitsandbytes, sentencepiece, protobuf, and flash-attn packages
 
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers import LlamaTokenizer, MixtralForCausalLM
-#import bitsandbytes, flash_attn
+from modelscope import AutoModelForCausalLM, AutoTokenizer
+
 
 class EndpointHandler:
     def __init__(self, path=""):
-        self.tokenizer = LlamaTokenizer.from_pretrained(path, trust_remote_code=True)
-        self.model = MixtralForCausalLM.from_pretrained(
-            path,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            load_in_8bit=False,
-            load_in_4bit=True,
-            #use_flash_attention_2=True
-            )
+        self.tokenizer =AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, device_map='auto')
+            
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         sys_prompt=data["prompt"]
         list=data["inputs"]
@@ -32,20 +22,8 @@ class EndpointHandler:
 
         #for chat in prompts:
             #print(chat)
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        outputs = self.model.generate(**inputs, max_new_tokens=20)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        encodeds = self.tokenizer.encode(prompt, return_tensors="pt")
-        model_inputs = encodeds.to(device)
-        self.model.to(device)
-        generated_ids = self.model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-        decoded = self.tokenizer.decode(generated_ids[0])
-        return decoded
-
-        """
-        encodeds = self.tokenizer.encode(prompt, return_tensors="pt")
-        model_inputs = encodeds.to(device)
-        self.model.to(device)
-        generated_ids = self.model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-        decoded = self.tokenizer.decode(generated_ids[0])
-        return decoded
-        """