File size: 3,067 Bytes
bcd03d2 | 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 | from typing import Dict, Any, List
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
class EndpointHandler:
def __init__(self, path=""):
"""Initialize the model and tokenizer.
Args:
path (str): Path to the model directory. Defaults to empty string.
"""
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = AutoModelForCausalLM.from_pretrained(
path or "merged",
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto"
)
self.tokenizer = AutoTokenizer.from_pretrained(path or "merged")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""Handle inference requests.
Args:
data (Dict[str, Any]): The input data. Can be in two formats:
1. Standard Hugging Face format:
{
"inputs": str,
"parameters": Dict[str, Any]
}
2. Custom format:
{
"instruction": str,
"input": str (optional),
"max_new_tokens": int (optional),
"temperature": float (optional)
}
Returns:
Dict[str, Any]: The model's response containing:
- response (str): The generated text
"""
# Handle standard Hugging Face format
if "inputs" in data:
instruction = data["inputs"]
parameters = data.get("parameters", {})
input_text = ""
max_new_tokens = parameters.get("max_new_tokens", 512)
temperature = parameters.get("temperature", 0.7)
# Handle custom format
else:
instruction = data.get("instruction", "")
input_text = data.get("input", "")
max_new_tokens = data.get("max_new_tokens", 512)
temperature = data.get("temperature", 0.7)
# Create prompt
prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{instruction}"""
if input_text:
prompt += f"""
### Input:
{input_text}"""
prompt += """
### Response:"""
# Generate response
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
outputs = self.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
do_sample=True,
pad_token_id=self.tokenizer.eos_token_id
)
# Decode and extract response
full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
response = full_response.split("### Response:")[-1].strip()
return {"response": response} |