fusion_text_detector.py

import torch
import torch.nn as nn
from transformers import BertTokenizer, BertForSequenceClassification
from perplexity_score import calculate_perplexity


class FusionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(2, 1)  # BERT score + Perplexity score

    def forward(self, bert_score, perplexity_score):
        x = torch.tensor([[bert_score, perplexity_score]], dtype=torch.float)
        return torch.sigmoid(self.fc(x))



tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
bert_model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased", num_labels=2
)
bert_model.eval()

fusion_model = FusionModel()


def predict_ai_probability(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs = bert_model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)
        bert_ai_score = probs[0][1].item()

    perplexity = calculate_perplexity(text)
    perplexity_norm = min(perplexity / 100, 1.0)

    final_score = fusion_model(bert_ai_score, perplexity_norm)
    return final_score.item()


sample = "This essay discusses the importance of artificial intelligence."
print("AI Generated Probability:", predict_ai_probability(sample))