File size: 2,992 Bytes
854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b 854779e 38c558b | 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 | # handler.py
from __future__ import annotations
import os
from typing import Any, Dict, List, Union
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
GEN_KW = {
"temperature": float(os.getenv("GEN_TEMPERATURE", "0.2")),
"do_sample": os.getenv("GEN_DO_SAMPLE", "false").lower() == "true",
"max_new_tokens": int(os.getenv("GEN_MAX_NEW_TOKENS", "512")),
"repetition_penalty": float(os.getenv("GEN_REP_PENALTY", "1.0")),
}
PROMPT_PREFIX = (
"Преобразуй дореформенный русский текст в современную орфографию, "
"сохранив смысл и пунктуацию. Верни только преобразованный текст.\n\nТекст:\n"
)
PROMPT_SUFFIX = "\n\nСовременный вариант:"
def _to_list(x: Union[str, List[str]]) -> List[str]:
if isinstance(x, list):
return [str(t) for t in x]
return [str(x)]
class EndpointHandler:
def __init__(self, model_dir: str):
# model_dir is the local path downloaded by the endpoint
self.device = "cuda" if torch.cuda.is_available() else "cpu"
# IMPORTANT: trust_remote_code to load custom arch "gpt_oss"
self.tokenizer = AutoTokenizer.from_pretrained(
model_dir, use_fast=True, trust_remote_code=True
)
self.model = AutoModelForCausalLM.from_pretrained(
model_dir,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto" if torch.cuda.is_available() else None,
trust_remote_code=True,
)
if not torch.cuda.is_available():
# For safety on CPU, disable KV cache to reduce RAM spikes
self.model.config.use_cache = False
self.model.eval()
def _prepare_inputs(self, texts: List[str]) -> Dict[str, Any]:
prompts = [f"{PROMPT_PREFIX}{t}{PROMPT_SUFFIX}" for t in texts]
toks = self.tokenizer(
prompts, return_tensors="pt", padding=True, truncation=True
)
return {k: v.to(self.model.device) for k, v in toks.items()}
@torch.inference_mode()
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
"""
Accepts { "inputs": "text" } or { "inputs": ["t1","t2",...] }
Returns [{ "generated_text": "..." }, ...]
"""
if "inputs" not in data:
return [{"error": "missing 'inputs'"}]
texts = _to_list(data["inputs"])
inputs = self._prepare_inputs(texts)
out = self.model.generate(**inputs, **GEN_KW)
results: List[Dict[str, str]] = []
# decode only the newly generated part
for i, seq in enumerate(out):
inp_len = inputs["input_ids"][i].shape[-1]
gen_part = seq[inp_len:]
text = self.tokenizer.decode(gen_part, skip_special_tokens=True).strip()
results.append({"generated_text": text})
return results
|