File size: 1,527 Bytes
383292f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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:
""" |