from modelscope import AutoModelForCausalLM, AutoTokenizer


class EndpointHandler:
    def __init__(self, path=""):
        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"]
        prompt=f"<|im_start|>system\n{sys_prompt}.<|im_end|>\n"
        for item in list:
            if item["role"]=="assistant":
                content=item["content"]
                prompt+=f"<|im_start|>assistant\n{content}<|im_end|>\n"
            else:
                content=item["content"]
                prompt+=f"<|im_start|>user\n{content}<|im_end|>\n"
        prompt+="<|im_start|>assistant\n"

        #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)