File size: 3,583 Bytes
564cdc6 |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import json
import logging
import torch
from typing import List
from typing import Dict, Any
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria
import torch
class MyStoppingCriteria(StoppingCriteria):
def __init__(self, target_sequence, prompt, tokenizer):
self.target_sequence = target_sequence
self.prompt = prompt
self.tokenizer = tokenizer
def __call__(self, input_ids, scores, **kwargs):
# Get the generated text as a string
generated_text = self.tokenizer.decode(input_ids[0])
generated_text = generated_text.replace(self.prompt, '')
# Check if the target sequence appears in the generated text
if self.target_sequence in generated_text:
return True # Stop generation
return False # Continue generation
def __len__(self):
return 1
def __iter__(self):
yield self
class EndpointHandler:
def __init__(self, model_dir=""):
# load model and processor from path
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForCausalLM.from_pretrained(model_dir, load_in_4bit=True, device_map="auto")
self.template = {
"prompt_input": """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n""",
"prompt_no_input": """Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n""",
"response_split": """### Response:"""
}
self.instruction = """Extract the start and end sequences for the categories 'personal information', 'work experience', 'education' and 'skills' from the following text in dictionary form"""
if torch.cuda.is_available():
self.device = "cuda"
else:
self.device = "cpu"
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
"""
Args:
data (:dict:):
The payload with the text prompt and generation parameters.
"""
# process input
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", None)
res = self.template["prompt_input"].format(
instruction=self.instruction, input=input
)
messages = [
{"role": "user", "content": res},
]
input_ids = self.tokenizer.apply_chat_template(
messages, truncation=True, add_generation_prompt=True, return_tensors="pt"
).input_ids
input_ids = input_ids.to(self.device)
# pass inputs with all kwargs in data
if parameters is not None:
outputs = self.model.generate(
input_ids=input_ids,
stopping_criteria=MyStoppingCriteria("</s>", inputs, self.tokenizer),
**parameters)
else:
outputs = self.model.generate(
input_ids=input_ids, max_new_tokens=32,
stopping_criteria=MyStoppingCriteria("</s>", inputs, self.tokenizer)
)
# postprocess the prediction
prediction = self.tokenizer.decode(outputs[0][input_ids.shape[1]:]) #, skip_special_tokens=True)
prediction = prediction.split("</s>")[0]
# TODO: add processing of the LLM output
return [{"generated_text": prediction}] |