File size: 4,338 Bytes
aa94a30 05f5ab0 aa94a30 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """
Custom handler for HuggingFace Inference Endpoints.
This handler loads the base model with the LoRA adapter for inference.
"""
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
# Configuration
BASE_MODEL = "Qwen/Qwen3-Coder-Next"
SYSTEM_PROMPT = """You are a Univer API expert. Generate JavaScript code to perform spreadsheet operations using the Univer API.
The code should be an async arrow function that uses the univerAPI global object.
IMPORTANT PRINCIPLES:
1. ALWAYS use spreadsheet formulas (setFormula) for calculations like sum, average, max, min, count, etc.
- Use setFormula('=SUM(...)'), setFormula('=AVERAGE(...)'), etc.
- Do NOT use JavaScript calculations (reduce, Math.max, etc.) for these operations
- Formulas are preferred because they update dynamically when data changes
2. Use the correct API for the task:
- To delete rows: use deleteRows(), NOT clear()
- To insert rows: use insertRows(), NOT manual approaches
3. ONLY use documented Univer API methods. Do NOT invent or guess method names."""
class EndpointHandler:
def __init__(self, path: str = ""):
"""
Initialize the handler by loading the model and tokenizer.
Args:
path: Path to the model directory (contains adapter files)
"""
# Load tokenizer from base model to avoid tokenizer_config compatibility issues.
# The adapter repo's tokenizer_config.json may have extra_special_tokens as a list,
# which causes AttributeError with some transformers versions.
print(f"Loading tokenizer from base model: {BASE_MODEL}...")
self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
print(f"Loading base model: {BASE_MODEL}...")
# Use 8-bit quantization for memory efficiency
bnb_config = BitsAndBytesConfig(
load_in_8bit=True,
llm_int8_enable_fp32_cpu_offload=False,
)
self.model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
print(f"Loading LoRA adapter from {path}...")
self.model = PeftModel.from_pretrained(self.model, path)
self.model.eval()
print("Model loaded successfully!")
def __call__(self, data: dict) -> dict:
"""
Handle inference requests.
Args:
data: Request data with 'inputs' key containing the instruction
Returns:
Generated code response
"""
# Extract input
inputs = data.get("inputs", data.get("input", ""))
parameters = data.get("parameters", {})
# Build messages
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": inputs},
]
# Apply chat template
prompt = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
# Tokenize
input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to(self.model.device)
# Generation parameters
max_new_tokens = parameters.get("max_new_tokens", 1024)
temperature = parameters.get("temperature", 0.1)
top_p = parameters.get("top_p", 0.9)
do_sample = parameters.get("do_sample", temperature > 0)
# Generate
with torch.no_grad():
outputs = self.model.generate(
input_ids,
max_new_tokens=max_new_tokens,
temperature=temperature if do_sample else None,
top_p=top_p if do_sample else None,
do_sample=do_sample,
pad_token_id=self.tokenizer.pad_token_id,
eos_token_id=self.tokenizer.eos_token_id,
)
# Decode response (only new tokens)
response = self.tokenizer.decode(
outputs[0][input_ids.shape[1] :],
skip_special_tokens=True,
)
return {"generated_text": response.strip()}
|