app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import re
import numpy as np
from tokenizers.normalizers import Sequence, Replace, Strip
from tokenizers import Regex
import matplotlib.pyplot as plt

# --- Setup and Model Loading ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Model paths and URLs
model1_path = "modernbert.bin"
model2_url = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed12"
model3_url = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed22"

tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")

def load_model(path_or_url, num_labels=41):
    model = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=num_labels)
    if path_or_url.startswith("http"):
        state_dict = torch.hub.load_state_dict_from_url(path_or_url, map_location=device)
    else:
        state_dict = torch.load(path_or_url, map_location=device)
    model.load_state_dict(state_dict)
    return model.to(device).eval()

# Initializing the ensemble
print("Loading Ensemble Models...")
model_1 = load_model(model1_path)
model_2 = load_model(model2_url)
model_3 = load_model(model3_url)
print("Models Loaded Successfully.")

# --- Text Preprocessing ---
def clean_text(text: str) -> str:
    text = re.sub(r'\s{2,}', ' ', text)
    text = re.sub(r'\s+([,.;:?!])', r'\1', text)
    return text

newline_to_space = Replace(Regex(r'\s*\n\s*'), " ")
join_hyphen_break = Replace(Regex(r'(\w+)[--]\s*\n\s*(\w+)'), r"\1\2")
tokenizer.backend_tokenizer.normalizer = Sequence([
    tokenizer.backend_tokenizer.normalizer,
    join_hyphen_break,
    newline_to_space,
    Strip()
])

# --- BRAND NAME MAPPING ---
# Mapping the architectural traces directly to consumer brands
openai_indices = [8, 18, 19, 20, 21, 39, 40]
# Updated Mapping:
modern_grouping = {
    "ChatGPT (OpenAI)": openai_indices, 
    "Meta AI (Llama)": [0, 1, 2, 3, 25, 26], 
    "Gemini (Google)": [11, 12, 13, 14, 15, 16, 17], 
    "Grok (xAI)": [4, 22, 23], # Removed 27 from here
    "Claude / Other AI": [5, 6, 7, 9, 10, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38] # Added 27 here
}

# --- STYLOMETRIC ANALYSIS ---
def calculate_stylometrics(text):
    sentences = re.split(r'[.!?]+', text)
    words = re.findall(r'\b\w+\b', text.lower())
    
    if not words or len(sentences) < 2:
        return 0, 0, 0, 0 
    
    lengths = [len(s.split()) for s in sentences if len(s.strip()) > 0]
    variance = np.var(lengths) if lengths else 0
    
    unique_words = set(words)
    ttr = len(unique_words) / len(words)
    
    human_boost = 0
    ai_boost = 0
    
    if variance > 50: 
        human_boost += 0.15
    elif variance < 20:
        ai_boost += 0.15
        
    if ttr > 0.65:
        human_boost += 0.05
    elif ttr < 0.45:
        ai_boost += 0.05
        
    return human_boost, ai_boost, variance, ttr

# --- Core Classification Logic ---
def classify_text(text):
    cleaned_text = clean_text(text)
    word_count = len(cleaned_text.split())
    if word_count < 30:
        return "Please enter at least 30 words for an accurate analysis.", None, "Insufficient data."

    inputs = tokenizer(cleaned_text, return_tensors="pt", truncation=True, padding=True).to(device)

    with torch.no_grad():
        logits_1 = model_1(**inputs).logits
        logits_2 = model_2(**inputs).logits
        logits_3 = model_3(**inputs).logits

        s1, s2, s3 = torch.softmax(logits_1, 1), torch.softmax(logits_2, 1), torch.softmax(logits_3, 1)
        avg_probs = (s1 + s2 + s3) / 3
        probs = avg_probs[0]

    # --- BALANCED RLHF SUPPRESSOR ---
    pre_suppression_openai = sum([probs[i].item() for i in openai_indices])
    
    # Soften the OpenAI dominance just enough to let Gemini and Grok traces show through
    for idx in openai_indices:
        probs[idx] = probs[idx] * 0.15  
    
    total_prob = probs.sum().item()
    probs = probs / total_prob
    # -------------------------------------

    base_human_prob = probs[24].item()
    
    ai_probs_only = probs.clone()
    ai_probs_only[24] = 0 
    base_ai_total_prob = ai_probs_only.sum().item()

    # Apply Stylometrics
    human_modifier, ai_modifier, variance, ttr = calculate_stylometrics(cleaned_text)
    
    adjusted_human_prob = max(0, base_human_prob + human_modifier)
    adjusted_ai_prob = max(0, base_ai_total_prob + ai_modifier)

    total_sum = adjusted_human_prob + adjusted_ai_prob
    human_pct = (adjusted_human_prob / total_sum) * 100
    ai_pct = (adjusted_ai_prob / total_sum) * 100

    # Calculate Brand Scores
    modern_scores = {}
    for modern_name, indices in modern_grouping.items():
        family_sum = sum([probs[i].item() for i in indices])
        modern_scores[modern_name] = family_sum

    total_ai_modern_score = sum(modern_scores.values())
    if total_ai_modern_score > 0:
        modern_scores_pct = {k: (v / total_ai_modern_score) * 100 for k, v in modern_scores.items()}
    else:
        modern_scores_pct = {k: 0 for k in modern_scores.keys()}

    sorted_families = sorted(modern_scores_pct.items(), key=lambda item: item[1], reverse=True)
    top_family, top_score = sorted_families[0]
    second_family, second_score = sorted_families[1]

    # --- Generate Diagnostics Report ---
    diagnostics = f"""
    **Lexical Analysis:**
    * Word Count: `{word_count}`
    * Sentence Variance (Burstiness): `{variance:.2f}` *(>50 favors Human, <20 favors AI)*
    * Type-Token Ratio (Richness): `{ttr:.2f}` *(>0.65 favors Human, <0.45 favors AI)*
    
    **Algorithmic Adjustments:**
    * Stylometric Modifiers Applied: `Human +{human_modifier:.2f} | AI +{ai_modifier:.2f}`
    * Raw RLHF Signal Detected (Pre-Suppression): `{pre_suppression_openai * 100:.1f}%`
    """

    # --- Construct Output ---
    if human_pct > ai_pct:
        result_message = (
            f"### Result: <span class='highlight-human'>**{human_pct:.2f}% Human written**</span>\n\n"
            "The content matches human writing patterns in both semantic structure and lexical variance."
        )
        
        fig, ax = plt.subplots(figsize=(8, 4))
        bars = ax.bar(['Human', 'AI'], [human_pct, ai_pct], color=['#4CAF50', '#FF5733'], alpha=0.8)
        ax.set_ylabel('Probability (%)')
        ax.set_title('Detection Probability')
        ax.set_ylim(0, 110)
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height + 2, f'{height:.1f}%', ha='center', fontweight='bold')
            
    else:
        result_message = (
            f"### Result: <span class='highlight-ai'>**{ai_pct:.2f}% AI generated**</span>\n\n"
            f"**Likely Source:** `{top_family}` ({top_score:.1f}% Match)\n"
            f"**Secondary Match:** `{second_family}` ({second_score:.1f}% Match)\n"
        )
        
        # Clean up names for the graph (e.g., "ChatGPT (OpenAI)" -> "ChatGPT")
        fig, ax = plt.subplots(figsize=(10, 5))
        families = [f[0].split('(')[0].strip() if '(' in f[0] else f[0] for f in sorted_families] 
        scores = [f[1] for f in sorted_families]
        
        bars = ax.bar(families, scores, color=['#10a37f', '#4285f4', '#000000', '#0668E1', '#D3D3D3'], alpha=0.8)
        ax.set_ylabel('Signal Strength within AI (%)')
        ax.set_title('AI Brand Attribution Breakdown')
        ax.set_ylim(0, max(scores) + 15 if max(scores) > 0 else 100)
        
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height + 1, f'{height:.1f}%', ha='center', fontweight='bold')

    plt.tight_layout()
    return result_message, fig, diagnostics

# --- Gradio UI Layout ---
with gr.Blocks(css="""
    .highlight-human { color: #4CAF50; font-weight: bold; background: rgba(76, 175, 80, 0.1); padding: 5px; border-radius: 5px; }
    .highlight-ai { color: #FF5733; font-weight: bold; background: rgba(255, 87, 51, 0.1); padding: 5px; border-radius: 5px; }
    #output-container { text-align: center; padding: 20px; }
""") as iface:
    gr.Markdown("# AI Text Detector & Brand Identifier")
    gr.Markdown("Detects AI-generated text and maps structural traces to major consumer models like **ChatGPT, Gemini, Grok, and Meta AI**.")
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Input Text (Min 30 Words required for mathematical accuracy)", 
                placeholder="Paste your English text here to begin analysis...", 
                lines=12
            )
        with gr.Column(scale=1):
            result_output = gr.Markdown("Analysis results will appear here.", elem_id="output-container")
            plot_output = gr.Plot()
            
            with gr.Accordion("Technical Diagnostics (Under the Hood)", open=False):
                diagnostics_output = gr.Markdown("Awaiting input to generate metrics...")

    text_input.change(
        classify_text, 
        inputs=text_input, 
        outputs=[result_output, plot_output, diagnostics_output]
    )

    gr.Markdown("---")
    gr.Markdown("**Engineered by SzegedAI**")

if __name__ == "__main__":
    iface.launch(share=True)