src/models/model.py - master - Walmet/CTF

src/models/model.py
61 строка2 KB
Maksim Goncharov
MVP
14 дек 2025, 23:49
14 дек 2025, 23:49a4c2b1d
import torch
import math
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from typing import Tuple

class QAModel:
    def __init__(self, model_name: str = 'google/flan-t5-base'):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.eval()

    def predict(self, text: str, prompt: str) -> Tuple[str, float]:
        full_input = f"question: {prompt} context: {text}"
        device = next(self.model.parameters()).device

        inputs = self.tokenizer(
            full_input,
            return_tensors="pt",
            max_length=512,
            truncation=True
        ).to(device)

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=20,
                output_scores=True,
                return_dict_in_generate=True,
                num_beams=3,
                do_sample=False
            )

        generated_ids = outputs.sequences[0]

        log_probs = []
        for i, scores_step in enumerate(outputs.scores):
            log_probs_step = F.log_softmax(scores_step, dim=-1)
            if (i + 1) < generated_ids.size(0):
                predicted_token_id = generated_ids[i + 1]
                log_prob = log_probs_step[0, predicted_token_id].item()
                log_probs.append(log_prob)

        if not log_probs:
            probability = 0.01
        else:
            average_log_score = torch.tensor(log_probs).mean().item()
            probability = max(0.01, min(1.0, math.exp(average_log_score)))

        generated_answer = self.tokenizer.decode(
            generated_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True
        ).strip()

        if generated_answer == "":
            probability = 0.01
        elif len(generated_answer.split()) > 10:
            probability *= 0.7

        return generated_answer, probability