import os
import torch
import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Change to your Hugging Face model repo ID
MODEL_ID = "SharvNey/capstone_project"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

def classify_text(text):
    if not text.strip():
        return {"🧑 Human-Written": 0.0, "🤖 AI-Generated": 0.0}
    enc = tokenizer(text, truncation=True, padding=True, max_length=256, return_tensors="pt")
    enc = {k: v.to(device) for k,v in enc.items()}
    with torch.no_grad():
        out = model(**enc)
        probs = torch.nn.functional.softmax(out.logits, dim=-1).cpu().numpy()[0]
    return {"🧑 Human-Written": float(probs[0]), "🤖 AI-Generated": float(probs[1])}

demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=8, placeholder="Paste text here..."),
    outputs=gr.Label(num_top_classes=2),
    title="🤖 AI vs Human Text Classifier",
    description="Fine-tuned RoBERTa model that detects whether text is Human-written 🧑 or AI-generated 🤖"
)

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