instruction-gen/

action-llama3-sft

class ActionGenerator:
    def __init__(self):
        self.base_model_path = "Meta-Llama-3-8B-Instruct-hf"
        self.adaptor_path = "THIS ADAPTER"
        self.device = torch.device("cuda")
        self.model, self.tokenizer = self.load_model_adapter()

    def load_model_adapter(self):
        model_tmp = AutoModelForCausalLM.from_pretrained(
            self.base_model_path,
            return_dict=True,
            torch_dtype=torch.float16,
        )
        tokenizer = AutoTokenizer.from_pretrained(self.base_model_path)
        tokenizer.pad_token = tokenizer.eos_token
        model = PeftModel.from_pretrained(model_tmp, self.adaptor_path)
        model = model.merge_and_unload()
        model = model.to(self.device)
        model.eval()
        return model, tokenizer

    def infer(self, messages):
        with torch.no_grad():
            input_ids = self.tokenizer.apply_chat_template(messages,add_generation_prompt=True,return_tensors="pt",padding=False,truncation=True,max_length=2048).to(self.device)
            terminators = [self.tokenizer.eos_token_id,self.tokenizer.convert_tokens_to_ids("<|eot_id|>")]
            outputs = self.model.generate(
                input_ids,
                max_new_tokens=256,
                eos_token_id=terminators,
                pad_token_id=self.tokenizer.pad_token_id,
                do_sample=True,
                temperature=0.8,
                top_p=0.9,
            )
            generated_ids = outputs[0][input_ids.shape[1]:]
            response = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
            torch.cuda.empty_cache()
            return response