import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model_path = "modernbert.bin"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
model = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device)
model.eval()

label_mapping = {
    0: '13B', 1: '30B', 2: '65B', 3: '7B', 4: 'GLM130B', 5: 'bloom_7b',
    6: 'bloomz', 7: 'cohere', 8: 'davinci', 9: 'dolly', 10: 'dolly-v2-12b',
    11: 'flan_t5_base', 12: 'flan_t5_large', 13: 'flan_t5_small', 
    14: 'flan_t5_xl', 15: 'flan_t5_xxl', 16: 'gemma-7b-it', 17: 'gemma2-9b-it',
    18: 'gpt-3.5-turbo', 19: 'gpt-35', 20: 'gpt4', 21: 'gpt4o', 
    22: 'gpt_j', 23: 'gpt_neox', 24: 'human', 25: 'llama3-70b', 26: 'llama3-8b',
    27: 'mixtral-8x7b', 28: 'opt_1.3b', 29: 'opt_125m', 30: 'opt_13b',
    31: 'opt_2.7b', 32: 'opt_30b', 33: 'opt_350m', 34: 'opt_6.7b',
    35: 'opt_iml_30b', 36: 'opt_iml_max_1.3b', 37: 't0_11b', 38: 't0_3b',
    39: 'text-davinci-002', 40: 'text-davinci-003'
}

def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    inputs = {key: value.to(device) for key, value in inputs.items()}
    
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=1)[0]
        predicted_class = torch.argmax(probabilities).item()
        confidence = probabilities[predicted_class].item() * 100  # Confidence as %

        if predicted_class == 24:
            prediction_label = f"✅ - The text is **{confidence:.2f}%** likely **Human written**."
            model_info = ""  
        else:
            prediction_label = f"🤖 - The text is **{confidence:.2f}%** likely **AI generated**."
            model_info = f"**Identified AI Model:** {label_mapping[predicted_class]}"

    result_message = f"**Result:**\n\n{prediction_label}"
    if model_info:
        result_message += f"\n\n{model_info}"
        
    return result_message

title = "AI Text Detector"
description = """
**AI detection tool by SzegedAI**

Detect AI-generated texts with precision using the new **ModernBERT** model, fine-tuned for machine-generated text detection, and capable of identifying 40 different models.

- **🤖 Identify AI Models**: Reveals which LLM generated the text if detected as AI.
- **✅ Human Verification**: Marks confidently human-written text with a green checkmark.

**Press the button below to classify your text!**


"""

iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(
        label="Enter Text for Analysis", 
        placeholder="Type or paste your content here...",
        lines=5,
        elem_id="text_input_box"
    ),
    outputs=gr.Markdown(elem_id="result_output_box"),
    title=title,
    description=description,
    allow_flagging="never", 
    live=False,  
    css="""
    #text_input_box {
        border-radius: 10px;
        border: 2px solid #4CAF50;
        font-size: 18px;
        padding: 10px;
    }
    #result_output_box {
        border-radius: 10px;
        border: 2px solid #4CAF50;
        font-size: 18px;
        padding: 15px;
        background-color: #2E2E3F;
        margin-top: 20px;
    }
    body {
        background: #1E1E2F;
        color: #E1E1E6;
        font-family: 'Aptos', sans-serif;
        padding: 20px;
    }
    .gradio-container {
        border: 2px solid #4CAF50;
        border-radius: 15px;
        padding: 25px;
        box-shadow: 0px 0px 20px rgba(0,255,0,0.6);
    }
    h1, h2 {
        text-align: center;
        font-size: 32px;
        font-weight: bold;
        margin-bottom: 20px;
    }
    """
)

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