| import os
|
| import json
|
| import torch
|
| import numpy as np
|
| import torch.nn.functional as F
|
| try:
|
| from .model import SimpleMCQModel
|
| from .tokenizer import MCQTokenizer
|
| except ImportError:
|
| from model import SimpleMCQModel
|
| from tokenizer import MCQTokenizer
|
|
|
| class MCQInferencePipeline:
|
| def __init__(self, model, tokenizer, label_map, max_len=128, device='cpu'):
|
| self.model = model.to(device)
|
| self.model.eval()
|
| self.tokenizer = tokenizer
|
| self.label_map = label_map
|
| self.max_len = max_len
|
| self.device = device
|
|
|
| @classmethod
|
| def load_from_dir(cls, model_dir, device='cpu'):
|
| with open(os.path.join(model_dir, 'config.json'), 'r', encoding='utf-8') as f:
|
| config = json.load(f)
|
|
|
| with open(os.path.join(model_dir, 'label_mapping.json'), 'r', encoding='utf-8') as f:
|
| raw_labels = json.load(f)
|
| label_map = {int(k): v for k, v in raw_labels.items()}
|
|
|
| tokenizer = MCQTokenizer.load_vocab(os.path.join(model_dir, 'vocab.json'), max_len=config.get('max_length', 128))
|
| model = SimpleMCQModel(vocab_size=config['vocab_size'], embed_dim=config['embedding_dim'], hidden_dim=config['hidden_dim'])
|
| model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pt'), map_location=device))
|
|
|
| return cls(model=model, tokenizer=tokenizer, label_map=label_map, max_len=config.get('max_length', 128), device=device)
|
|
|
| def predict(self, prompt, options):
|
| """
|
| options: list of 5 text options [optA, optB, optC, optD, optE]
|
| """
|
| option_tensors = []
|
| for opt_text in options:
|
| combined_text = str(prompt) + " " + str(opt_text)
|
| tokens = self.tokenizer.tokenize(combined_text)
|
| option_tensors.append(tokens)
|
|
|
|
|
| x = torch.tensor([option_tensors], dtype=torch.long, device=self.device)
|
| with torch.no_grad():
|
| logits = self.model(x)
|
| probs = F.softmax(logits, dim=1).squeeze(0).cpu().numpy()
|
|
|
| sorted_indices = np.argsort(probs)[::-1]
|
| top1_idx = sorted_indices[0]
|
| top1_label = self.label_map[top1_idx]
|
| top3_labels = [self.label_map[i] for i in sorted_indices[:3]]
|
|
|
| confidence_scores = {self.label_map[i]: float(probs[i]) for i in range(len(options))}
|
|
|
| return {
|
| "top1_label": top1_label,
|
| "top1_option_text": options[top1_idx],
|
| "top1_confidence": float(probs[top1_idx]),
|
| "top3_labels": top3_labels,
|
| "top3_str": " ".join(top3_labels),
|
| "confidence_scores": confidence_scores
|
| }
|
|
|
