初始化项目,由ModelHub XC社区提供模型
Model: Dilana/Llama-3.2-1B-Adaptive-RAG-v3 Source: Original Platform
This commit is contained in:
40
handler.py
Normal file
40
handler.py
Normal file
@@ -0,0 +1,40 @@
|
||||
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
|
||||
Reference in New Issue
Block a user