| import torch | |
| from peft import AutoModelForCausalLM | |
| from transformers import AutoTokenizer | |
| from cog import BasePredictor, Input | |
| class Predictor(BasePredictor): | |
| def setup(self): | |
| model_id = 'pbevan11/llama-3-8b-ocr-correction' | |
| self.model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True) | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| self.tokenizer.pad_token = self.tokenizer.eos_token | |
| def predict(self, instruction: str = Input(description="Instruction for the model"), | |
| inp: str = Input(description="Input text to correct")) -> str: | |
| prompt = self.create_prompt(instruction, inp) | |
| input_ids = self.tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cuda() | |
| out_ids = self.model.generate(input_ids=input_ids, max_new_tokens=5000, do_sample=False) | |
| full_output = self.tokenizer.batch_decode(out_ids.detach().cpu().numpy(), skip_special_tokens=True)[0] | |
| response_start = full_output.find("### Response:") | |
| if response_start != -1: | |
| return full_output[response_start + len("### Response:"):] | |
| else: | |
| return full_output[len(prompt):] | |
| def create_prompt(self, instruction, inp): | |
| return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. | |
| ### Instruction: | |
| {instruction} | |
| ### Input: | |
| {inp} | |
| ### Response: | |
| """ |