|
|
from typing import Dict, List, Any |
|
|
import json |
|
|
import torch |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer |
|
|
|
|
|
class EndpointHandler(): |
|
|
def __init__(self, path=""): |
|
|
self.tokenizer = None |
|
|
self.model = None |
|
|
self.device = None |
|
|
self.load_model(path) |
|
|
|
|
|
def load_model(self, model_dir): |
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
model_path = model_dir |
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(self.device) |
|
|
self.model.eval() |
|
|
print(f"Tokenizer and Model loaded from: {model_path} to device: {self.device}") |
|
|
|
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
|
inputs = data.pop("inputs", data) |
|
|
print(f'get input {inputs}') |
|
|
if not inputs: |
|
|
raise ValueError("Input text is missing in the request. Please provide 'inputs' or 'text' in your request.") |
|
|
|
|
|
history = json.loads(inputs) |
|
|
print(f'history is {history}') |
|
|
|
|
|
conversion = self.tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=False) |
|
|
encoding = self.tokenizer(conversion, return_tensors="pt").to(self.device) |
|
|
print(f'encoding success') |
|
|
with torch.no_grad(): |
|
|
output = self.model.generate( |
|
|
**encoding, |
|
|
max_new_tokens=1024, |
|
|
temperature=1.5, |
|
|
do_sample=True, |
|
|
pad_token_id=self.tokenizer.eos_token_id |
|
|
) |
|
|
print(f'output success') |
|
|
generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True) |
|
|
|
|
|
return [{"response": generated_text}] |
|
|
|
