import gradio as gr
from transformers import pipeline

# Load the model directly from your Hugging Face account
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 [ 0 ] to extract the dictionary from the nested list
    inner_results = results[ 0 ]
    
    return {res['label']: res['score'] for res in inner_results}

class_intro_html = "<center><h2>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()