40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
|
from typing import Dict, List, Any
|
||
|
|
import torch
|
||
|
|
from transformers import PreTrainedTokenizerFast, AutoModelForCausalLM, pipeline
|
||
|
|
|
||
|
|
class EndpointHandler():
|
||
|
|
def __init__(self, path=""):
|
||
|
|
# 1. Force the Fast Tokenizer to bypass the legacy TokenizersBackend bug
|
||
|
|
# Setting padding_side="left" is standard for Causal LMs
|
||
|
|
self.tokenizer = PreTrainedTokenizerFast.from_pretrained(
|
||
|
|
path,
|
||
|
|
padding_side="left"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Safely set pad token to eos_token to avoid warnings
|
||
|
|
if self.tokenizer.pad_token_id is None:
|
||
|
|
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
|
||
|
|
|
||
|
|
# 2. Load Model strictly for CPU
|
||
|
|
# Using bfloat16 optimizes RAM usage significantly on Sapphire Rapids
|
||
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
path,
|
||
|
|
torch_dtype=torch.bfloat16
|
||
|
|
)
|
||
|
|
|
||
|
|
# 3. Create Pipeline attached to self
|
||
|
|
self.pipeline = pipeline(
|
||
|
|
"text-generation",
|
||
|
|
model=self.model,
|
||
|
|
tokenizer=self.tokenizer,
|
||
|
|
return_full_text=False # Prevent echoing the prompt
|
||
|
|
)
|
||
|
|
|
||
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||
|
|
# Extract inputs and parameters safely
|
||
|
|
inputs = data.pop("inputs", data)
|
||
|
|
parameters = data.pop("parameters", {})
|
||
|
|
|
||
|
|
# Execute generation
|
||
|
|
prediction = self.pipeline(inputs, **parameters)
|
||
|
|
return prediction
|