| from typing import List, Tuple, Any | |
| import numpy as np | |
| from transformers import AutoTokenizer | |
| from bi_encoder import get_corpus, get_question | |
| import torch | |
| from model import CrossEncoderBert | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = CrossEncoderBert() | |
| model.model.resize_token_embeddings(len(model.tokenizer)) | |
| model.load_state_dict(torch.load('model/torch_model', map_location=torch.device(device))) | |
| model.tokenizer = AutoTokenizer.from_pretrained('model/tokenizer') | |
| model.to(device) | |
| def get_range_answers( | |
| context: str, | |
| question: str, | |
| num_answers: int = 5) -> list[str]: | |
| corpus = get_corpus(context, question) | |
| context_question = f'{context} [Cont_token] {question}' | |
| context_questions = [context_question] * len(corpus) | |
| tokenized_texts = model.tokenizer( | |
| context_questions, | |
| corpus, | |
| max_length=512, | |
| padding=True, | |
| truncation=True, | |
| return_tensors='pt' | |
| ).to(device) | |
| with torch.no_grad(): | |
| ce_scores = model(tokenized_texts['input_ids'], | |
| tokenized_texts['attention_mask']).squeeze(-1) | |
| ce_scores = torch.sigmoid(ce_scores) | |
| scores = ce_scores.cpu().numpy() | |
| scores_ix = np.argsort(scores)[::-1] | |
| best_answers = [] | |
| for idx in scores_ix[:num_answers]: | |
| best_answers.append((scores[idx], corpus[idx])) | |
| best_answers = [str(x[1]) for x in best_answers] | |
| return best_answers | |
| def get_best_answer( | |
| context: str, | |
| question: str | |
| ) -> str: | |
| return get_range_answers(context, question, 1) | |