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: **{human_pct:.2f}% Human written**\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: **{ai_pct:.2f}% AI generated**\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)