80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""
|
|
HuggingFace Inference Endpoint custom handler for sinllama-mcq-merged-2.0.
|
|
|
|
The deployed model (itsjorigo/sinllama-mcq-merged) is a fully merged
|
|
AutoModelForCausalLM — NOT a PEFT adapter stack. Load it directly.
|
|
|
|
The tokenizer must come from polyglots/SinLlama_v01 (trust_remote_code=True)
|
|
because that repo defines the custom TokenizersBackend class.
|
|
|
|
Request: {"inputs": {"passage": "<Sinhala text>", "entity_block": "<optional>"}}
|
|
Response: {"mcq": "ප්රශ්නය: ..."}
|
|
"""
|
|
|
|
import re
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
SINLLAMA_ID = "polyglots/SinLlama_v01"
|
|
|
|
PROMPT_TEMPLATE = (
|
|
"පහත ඉතිහාස ඡේදය කියවා, ඒ ගැන බහු-විකල්ප ප්රශ්නයක් සාදන්න.\n\n"
|
|
"ඡේදය: {passage}\n\n"
|
|
"{entity_block}"
|
|
"MCQ:"
|
|
)
|
|
|
|
|
|
class EndpointHandler:
|
|
def __init__(self, path=""):
|
|
# Tokenizer from SinLlama — defines the TokenizersBackend class
|
|
print("Loading tokenizer from SinLlama repo...", flush=True)
|
|
self.tokenizer = AutoTokenizer.from_pretrained(SINLLAMA_ID, trust_remote_code=True)
|
|
|
|
# Model loaded directly — itsjorigo/sinllama-mcq-merged is already fully merged
|
|
print(f"Loading merged model from {path!r}...", flush=True)
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
path,
|
|
torch_dtype=torch.float16,
|
|
device_map="auto",
|
|
low_cpu_mem_usage=True,
|
|
attn_implementation="sdpa",
|
|
)
|
|
self.model.eval()
|
|
print("EndpointHandler ready.", flush=True)
|
|
|
|
def __call__(self, data: dict) -> dict:
|
|
inputs = data.get("inputs", {})
|
|
passage = inputs.get("passage", "")
|
|
entity_block = inputs.get("entity_block", "")
|
|
|
|
if not passage:
|
|
return {"error": 'No passage provided. Send {"inputs": {"passage": "..."}}'}
|
|
|
|
prompt = PROMPT_TEMPLATE.format(
|
|
passage=passage.strip(),
|
|
entity_block=entity_block,
|
|
)
|
|
|
|
enc = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
|
|
|
|
with torch.no_grad():
|
|
out = self.model.generate(
|
|
**enc,
|
|
max_new_tokens=280,
|
|
temperature=0.7,
|
|
do_sample=True,
|
|
repetition_penalty=1.1,
|
|
eos_token_id=self.tokenizer.eos_token_id,
|
|
pad_token_id=self.tokenizer.pad_token_id,
|
|
)
|
|
|
|
new_ids = out[0][enc.input_ids.shape[1]:]
|
|
text = self.tokenizer.decode(new_ids, skip_special_tokens=True).strip()
|
|
|
|
# Ensure each option starts on its own line
|
|
for tag in ["A)", "B)", "C)", "D)", "නිවැරදි පිළිතුර:"]:
|
|
text = re.sub(rf"(?<!\n)({re.escape(tag)})", r"\n\1", text)
|
|
|
|
return {"mcq": text.strip()}
|