import gradio as gr
from transformers import pipeline

# Load your custom model
classifier = pipeline("text-classification", model="jay123jay/AI-Text-Origin-Classifier", truncation=True, max_length=512, top_k=None)

def detect_text(text):
    if not text.strip():
        return {}
    
    results = classifier(text)
    
    # Use.pop(0) to safely extract the inner list without using brackets
    inner_results = results.pop(0)
    
    return {res['label']: res['score'] for res in inner_results}

class_intro_html = "<center><h2>AI Text Origin & LLM Attribution Detector</h2></center>"

with gr.Blocks() as demo:
    gr.HTML(class_intro_html)
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(lines=8, placeholder="Paste the text here to analyze its origin...", label="Input Text")
            analyze_button = gr.Button("Analyze Text")
        with gr.Column():
            output_labels = gr.Label(label="Attribution Probabilities")
            
    analyze_button.click(fn=detect_text, inputs=input_text, outputs=output_labels)

demo.launch()