import torch
import torch.nn as nn
import gradio as gr
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer

# 1. ARCHITECTURE DEFINITION
class MultiSampleDropoutHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
        self.multi_dropout = nn.ModuleList([nn.Dropout(0.1 * (i+1)) for i in range(5)])

    def forward(self, features, **kwargs):
        x = features[:, 0, :] 
        x = torch.tanh(self.dense(x))
        return sum([self.out_proj(d(x)) for d in self.multi_dropout]) / 5

# 2. MODEL LOADER
def get_model():
    repo = "KnoxfireyNeurons/Roberta_Base"
    tokenizer = AutoTokenizer.from_pretrained(repo)
    config = AutoConfig.from_pretrained(repo)
    model = AutoModelForSequenceClassification.from_pretrained(repo, config=config)
    model.classifier = MultiSampleDropoutHead(config)
    model.eval()
    return model, tokenizer

model, tokenizer = get_model()

# 3. PREDICTION ENGINE (FIXED)
def predict(text):
    if not text or not text.strip():
        return {"Please enter a review": 0.0}
    
    inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=256)
    
    with torch.no_grad():
        outputs = model(**inputs)
        # FIX: We must extract .logits from the output object
        logits = outputs.logits
    
    # Apply softmax to the extracted logits
    probs = torch.nn.functional.softmax(logits, dim=-1).squeeze().tolist()
    
    return {
        "Original (Real)": probs[0],
        "Computer Generated (Fake)": probs[1]
    }

# 4. INTERFACE
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Review Content", lines=5),
    outputs=gr.Label(num_top_classes=2, label="AI Analysis"),
    title="🤖 Fake Review Detector",
    description="Analyze reviews for AI patterns using a SOTA RoBERTa model.",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()