| from transformers import AutoTokenizer, LlamaForCausalLM, LlamaConfig |
| import torch |
| import os |
| from safetensors.torch import load_file |
|
|
| TOKENIZER_NAME = "polyglots/Extended-Sinhala-LLaMA" |
| VOCAB_SIZE = 139336 |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| print("Loading tokenizer...") |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| TOKENIZER_NAME, |
| trust_remote_code=True, |
| ) |
| if self.tokenizer.pad_token is None: |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
|
|
| print("Loading config...") |
| config = LlamaConfig.from_pretrained(path) |
| config.vocab_size = VOCAB_SIZE |
|
|
| print("Building model at correct vocab size...") |
| self.model = LlamaForCausalLM(config) |
|
|
| print("Loading checkpoint shards...") |
| weights = {} |
| for f in sorted(os.listdir(path)): |
| if f.endswith(".safetensors"): |
| print(f" Shard: {f}") |
| weights.update(load_file(os.path.join(path, f))) |
|
|
| missing, unexpected = self.model.load_state_dict(weights, strict=False) |
| print(f" Missing: {len(missing)} | Unexpected: {len(unexpected)}") |
|
|
| self.model = self.model.to(torch.float16).to("cuda") |
| self.model.config.pad_token_id = self.tokenizer.eos_token_id |
| self.model.eval() |
| print(f"Ready! Vocab: {self.model.config.vocab_size:,}") |
|
|
| def __call__(self, data: dict) -> dict: |
| inputs = data.get("inputs", "") |
| params = data.get("parameters", {}) |
| max_tokens = params.get("max_new_tokens", 400) |
| temperature = params.get("temperature", 0.7) |
| top_p = params.get("top_p", 0.9) |
| rep_penalty = params.get("repetition_penalty", 1.1) |
|
|
| if not inputs: |
| return {"error": "No input provided. Use the inputs key."} |
|
|
| tokenized = self.tokenizer( |
| inputs, |
| return_tensors="pt", |
| truncation=True, |
| max_length=1024, |
| ).to(self.model.device) |
|
|
| with torch.no_grad(): |
| output_ids = self.model.generate( |
| **tokenized, |
| max_new_tokens=max_tokens, |
| temperature=temperature, |
| top_p=top_p, |
| repetition_penalty=rep_penalty, |
| do_sample=True, |
| pad_token_id=self.tokenizer.eos_token_id, |
| ) |
|
|
| new_tokens = output_ids[0][tokenized.input_ids.shape[1]:] |
| decoded = self.tokenizer.decode(new_tokens, skip_special_tokens=True) |
| return {"generated_text": decoded.strip()} |
|
|
