File size: 1,936 Bytes
c8b5fa1 |
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 |
from typing import Dict, List, Any
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftConfig, PeftModel
class EndpointHandler():
def __init__(self, path=""):
config = PeftConfig.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
# Load the Lora model
self.model = PeftModel.from_pretrained(model, path)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
prompt (:obj:`str`):
temperature (:obj:`float`, `optional`, defaults to 0.5):
eos_token_id (:obj:`int`, `optional`, defaults to tokenizer.eos_token_id):
early_stopping (:obj:`bool`, `optional`, defaults to `True`):
repetition_penalty (:obj:`float`, `optional`, defaults to 0.3):
Return:
A :obj:`str` : generated sequences
"""
# Get inputs
prompt = data.pop("prompt", None)
temperature = data.pop("temperature", 0.5)
eos_token_id = data.pop("eos_token_id", self.tokenizer.eos_token_id)
early_stopping = data.pop('early_stopping', True)
repetition_penalty = data.pop('repetition_penalty', 0.3)
max_new_tokens = data.pop('max_new_tokens', 100)
if prompt is None:
raise ValueError("No prompt provided.")
# Run prediction
inputs = self.tokenizer(prompt, return_tensors="pt")
prediction = self.model.generate(
**inputs,
temperature=temperature,
eos_token_id=eos_token_id,
early_stopping=early_stopping,
repetition_penalty=repetition_penalty,
max_new_tokens=max_new_tokens
)
return prediction
|