|
from typing import List, Tuple
|
|
import numpy as np
|
|
import pandas as pd
|
|
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, RobertaForSequenceClassification, RobertaTokenizer, ElectraModel, ElectraForCausalLM, GPT2Tokenizer, GPT2Model, GPT2LMHeadModel
|
|
import torch
|
|
import os
|
|
import json
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.metrics.pairwise import cosine_similarity
|
|
import gzip
|
|
from transformers import Text2TextGenerationPipeline
|
|
|
|
class EnsembleQAPipeline(Text2TextGenerationPipeline):
|
|
def __init__(self, model=None, tokenizer=None, framework="pt", **kwargs):
|
|
super().__init__(model=model, tokenizer=tokenizer, framework=framework)
|
|
self.quiz_bowl_model = QuizBowlModel()
|
|
|
|
def _forward(self, model_inputs, **generate_kwargs):
|
|
questions = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in model_inputs["input_ids"]]
|
|
results = self.quiz_bowl_model.guess_and_buzz(questions)
|
|
return results
|
|
|
|
def postprocess(self, model_outputs):
|
|
results = {}
|
|
for output in model_outputs:
|
|
guess_text = output[0]
|
|
confidence = output[1]
|
|
results = {'guess': guess_text, 'confidence': confidence}
|
|
return results
|
|
|
|
|
|
class QuizBowlModel:
|
|
def __init__(self):
|
|
model_configs = {
|
|
'flan-t5-large': {'model': 'google/flan-t5-large', 'tokenizer': 'google/flan-t5-large'},
|
|
'flan-t5-small': {'model': 'google/flan-t5-small', 'tokenizer': 'google/flan-t5-small'},
|
|
'flan-t5-base': {'model': 'google/flan-t5-base', 'tokenizer': 'google/flan-t5-base'}
|
|
}
|
|
self.models = {}
|
|
self.tokenizers = {}
|
|
self.load_models(model_configs)
|
|
|
|
def load_models(self, model_configs):
|
|
"""Load multiple models based on configuration."""
|
|
for model_name, config in model_configs.items():
|
|
tokenizer = AutoTokenizer.from_pretrained(config['tokenizer'])
|
|
model = AutoModelForSeq2SeqLM.from_pretrained(config['model'])
|
|
model.eval()
|
|
self.models[model_name] = model
|
|
self.tokenizers[model_name] = tokenizer
|
|
|
|
def guess_and_buzz(self, question_texts):
|
|
total_answers = []
|
|
for question in question_texts:
|
|
answers = self.generate_answers(question)
|
|
buzzed_answer, buzz_confidence = self.decide_to_buzz(answers)
|
|
total_answers.append((buzzed_answer, buzz_confidence))
|
|
return total_answers
|
|
|
|
def generate_answers(self, question):
|
|
raw_answers = []
|
|
for model_name, model in self.models.items():
|
|
tokenizer = self.tokenizers[model_name]
|
|
input_ids = tokenizer(question, return_tensors="pt", padding=True, truncation=True).input_ids
|
|
with torch.no_grad():
|
|
outputs = model.generate(input_ids, max_new_tokens=5, output_scores=True, return_dict_in_generate=True)
|
|
decoded_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
|
|
confidence_score = self.calculate_confidence(outputs.scores)
|
|
raw_answers.append((decoded_text, confidence_score))
|
|
return raw_answers
|
|
|
|
def calculate_confidence(self, scores):
|
|
if scores:
|
|
log_probs = [torch.nn.functional.log_softmax(score, dim=-1) for score in scores]
|
|
selected_log_probs = [log_probs[i][0, scores[i].argmax()].item() for i in range(len(log_probs))]
|
|
confidence_score = np.exp(np.mean(selected_log_probs))
|
|
else:
|
|
confidence_score = None
|
|
return confidence_score
|
|
|
|
def decide_to_buzz(self, answers):
|
|
cumulative_confidence = 0
|
|
for _, confidence in answers:
|
|
cumulative_confidence += confidence
|
|
if cumulative_confidence > 0.6:
|
|
return max(answers, key=lambda x: x[1])
|
|
return None, 0
|
|
|
