File size: 3,096 Bytes
09daf0b | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """
model3.py — Integration for BiLSTM Model.
"""
import logging
import torch
import os
from models.qa_model import QAModel
# Import vocab utilities and preprocess utilities
from utils.preprocess import tokenize
from utils.vocab import encode
logger = logging.getLogger(__name__)
model = None
vocab = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def init_model3():
global model, vocab
logger.info("[Model3] Initialising BiLSTM from qa_model.pth...")
# Assumes qa_model.pth is at the root of the backend directory
model_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "qa_model.pth")
if not os.path.exists(model_path):
logger.warning(f"[Model3] qa_model.pth not found at {model_path}! Model 3 inference will fail.")
return
try:
checkpoint = torch.load(model_path, map_location=device)
vocab = checkpoint["vocab"]
model = QAModel(len(vocab))
model.load_state_dict(checkpoint["model_state"])
model.to(device)
model.eval()
logger.info("[Model3] BiLSTM successfully loaded.")
except Exception as e:
logger.error(f"[Model3] Failed to load BiLSTM model: {e}")
def predict(context: str, question: str) -> dict:
"""Predict using the loaded BiLSTM."""
if model is None or vocab is None:
return {
"answer": "BiLSTM model weights (qa_model.pth) not found or failed to load. Please make sure the trained model is placed in the backend folder.",
"score": 0.0,
"model": "BiLSTM",
"model_id": "model3",
"error": True,
"stub": False,
}
try:
q_tokens = tokenize(question)
c_tokens = tokenize(context)
tokens = q_tokens + ["[SEP]"] + c_tokens
encoded = encode(tokens, vocab)
max_len = 300
if len(encoded) < max_len:
encoded += [0] * (max_len - len(encoded))
else:
encoded = encoded[:max_len]
x = torch.tensor(encoded).unsqueeze(0).to(device)
with torch.no_grad():
start_logits, end_logits = model(x)
start = torch.argmax(start_logits, dim=1).item()
end = torch.argmax(end_logits, dim=1).item()
if start > end or start >= len(tokens):
answer = "No answer found"
score = 0.0
else:
answer = " ".join(tokens[start:end+1])
# Extract basic score approximations from logits if needed, but returning dummy score for now.
score = 0.85
return {
"answer": answer,
"score": score,
"model": "BiLSTM",
"model_id": "model3",
"error": False,
}
except Exception as e:
logger.error(f"[Model3] Inference error: {e}")
return {
"answer": "Inference error occurred.",
"score": 0.0,
"model": "BiLSTM",
"model_id": "model3",
"error": True,
"stub": False,
}
|