parallelstudios/mpt-7b-instruct-parallel-colony-memory-importance-ft

1import warnings2import torch3from transformers import AutoModelForCausalLM, AutoTokenizer4from typing import Any, Dict5 6class EndpointHandler:7    INSTRUCTION_KEY = "### Instruction:"8    RESPONSE_KEY = "### Response:"9    END_KEY = "### End"10    INTRO_BLURB = "Below is an instruction that describes a task. Write a response that appropriately completes the request."11    PROMPT_FOR_GENERATION_FORMAT = """{intro}12    {instruction_key}13    {instruction}14    {response_key}15    """.format(16        intro=INTRO_BLURB,17        instruction_key=INSTRUCTION_KEY,18        instruction="{instruction}",19        response_key=RESPONSE_KEY,20    )21 22    def __init__(23        self,24        path,25        torch_dtype=torch.bfloat16,26        trust_remote_code=True,27    ) -> None:28        self.model = AutoModelForCausalLM.from_pretrained(29            path,30            torch_dtype=torch_dtype,31            trust_remote_code=trust_remote_code32        )33        tokenizer = AutoTokenizer.from_pretrained(34            "mosaicml/mpt-7b-instruct",35            trust_remote_code=trust_remote_code36        )37        if tokenizer.pad_token_id is None:38            warnings.warn(39                "pad_token_id is not set for the tokenizer. Using eos_token_id as pad_token_id."40            )41            tokenizer.pad_token = tokenizer.eos_token42 43        tokenizer.padding_side = "right"44        self.tokenizer = tokenizer45 46        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")47        self.model.eval()48        self.model.to(device=self.device, dtype=torch_dtype)49 50        self.generate_kwargs = {51            "temperature": 0.01,52            "top_p": 0.92,53            "top_k": 0,54            "max_new_tokens": 512,55            "use_cache": True,56            "do_sample": True,57            "eos_token_id": self.tokenizer.eos_token_id,58            "pad_token_id": self.tokenizer.pad_token_id,59            "repetition_penalty": 1.060        }61 62    def format_instruction(self, instruction):63        return self.PROMPT_FOR_GENERATION_FORMAT.format(instruction=instruction)64 65    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:66        # process input67        inputs = data.pop("inputs", data)68        parameters = data.pop("parameters", None)69 70        # preprocess71        s = self.format_instruction(instruction=inputs)72        input_ids = self.tokenizer(s, return_tensors="pt").input_ids.to(self.device)73        gkw = {**self.generate_kwargs, **parameters}74        # pass inputs with all kwargs in data  75        with torch.no_grad():76            output_ids = self.model.generate(input_ids, **gkw)77        # Slice the output_ids tensor to get only new tokens78        new_tokens = output_ids[0, len(input_ids[0]) :]79        output_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)80        return [{"generated_text": output_text}]