|
|
import os, sys, re |
|
|
from typing import Dict, List, Any, Union |
|
|
import torch |
|
|
|
|
|
REPO_ROOT = os.path.dirname(os.path.abspath(__file__)) |
|
|
if REPO_ROOT not in sys.path: |
|
|
sys.path.insert(0, REPO_ROOT) |
|
|
|
|
|
from model.model import Transformer |
|
|
from model.vocab.tokenizer import Tokenizer |
|
|
import config |
|
|
|
|
|
|
|
|
class EndpointHandler: |
|
|
def __init__(self, path: str = ""): |
|
|
self.base_dir = path or REPO_ROOT |
|
|
|
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
|
|
|
ckpt_path = os.path.join(self.base_dir, "epoch_10.pt") |
|
|
if not os.path.isfile(ckpt_path): |
|
|
raise FileNotFoundError(f"Missing checkpoint at: {ckpt_path}") |
|
|
|
|
|
self.model = Transformer().to(self.device) |
|
|
|
|
|
ckpt = torch.load(ckpt_path, map_location=self.device) |
|
|
|
|
|
if isinstance(ckpt, dict) and "state_dict" in ckpt: |
|
|
state_dict = ckpt["state_dict"] |
|
|
elif isinstance(ckpt, dict) and "model_state_dict" in ckpt: |
|
|
state_dict = ckpt["model_state_dict"] |
|
|
else: |
|
|
state_dict = ckpt |
|
|
|
|
|
self.model.load_state_dict(state_dict, strict=True) |
|
|
self.model.eval() |
|
|
|
|
|
|
|
|
token_path = os.path.join(self.base_dir, "tokenizer.model") |
|
|
if not os.path.isfile(token_path): |
|
|
raise FileNotFoundError(f"Missing tokenizer weights at: {token_path}") |
|
|
|
|
|
self.tokenizer = Tokenizer() |
|
|
self.tokenizer.load_weights(token_path) |
|
|
|
|
|
def _last_token_logits(self, model_out: torch.Tensor) -> torch.Tensor: |
|
|
if model_out.dim() == 3: |
|
|
return model_out[0, -1, :] |
|
|
if model_out.dim() == 2: |
|
|
return model_out[-1, :] |
|
|
raise ValueError(f"Unexpected model output shape: {tuple(model_out.shape)}") |
|
|
|
|
|
@torch.inference_mode() |
|
|
def _generate_one(self, prompt: str) -> str: |
|
|
encoded = torch.as_tensor( |
|
|
self.tokenizer.encode(prompt), |
|
|
dtype=torch.long, |
|
|
device=self.device, |
|
|
) |
|
|
|
|
|
if encoded.numel() == 0: |
|
|
return "AURELIUS: (No input processed)" |
|
|
|
|
|
currtoken = "" |
|
|
outputstring = "" |
|
|
countcheck = 0 |
|
|
|
|
|
while currtoken != "<END>" and countcheck < config.max_tokens: |
|
|
logits = self._last_token_logits(self.model(encoded)) |
|
|
|
|
|
if config.argmax: |
|
|
next_id = int(torch.argmax(logits).item()) |
|
|
else: |
|
|
probs = torch.softmax(logits / config.temperature, dim=-1) |
|
|
next_id = int(torch.multinomial(probs, num_samples=1).item()) |
|
|
|
|
|
currtoken = self.tokenizer.decode([next_id]).strip() |
|
|
|
|
|
if re.match(r"^[.,!?;:]", currtoken): |
|
|
if outputstring.endswith(" "): |
|
|
outputstring = outputstring[:-1] |
|
|
outputstring += currtoken + " " |
|
|
else: |
|
|
outputstring += currtoken + " " |
|
|
|
|
|
encoded = torch.cat( |
|
|
[encoded, torch.tensor([next_id], dtype=torch.long, device=self.device)], |
|
|
dim=0, |
|
|
) |
|
|
if encoded.numel() > config.max_seq_length: |
|
|
encoded = encoded[-config.max_seq_length :] |
|
|
|
|
|
countcheck += 1 |
|
|
|
|
|
text = re.sub("<BEGIN>", "\n\n", outputstring) |
|
|
text = re.sub("<END>", "\n\n", text) |
|
|
return "AURELIUS: " + text |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
|
inputs = data.get("inputs", data) |
|
|
|
|
|
if isinstance(inputs, dict): |
|
|
inputs = inputs.get("text", "") |
|
|
|
|
|
if isinstance(inputs, list): |
|
|
return [{"generated_text": self._generate_one(str(x))} for x in inputs] |
|
|
|
|
|
return [{"generated_text": self._generate_one(str(inputs))}] |
|
|
|
