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

# New multi-label model
MODEL_NAME = "SterlingWork/sdg-classifier-multilabel"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def predict(text):
    if not text or not text.strip():
        return {f"SDG {i}": 0.0 for i in range(1, 17)}
    
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding=True
    ).to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits[0]
        probs = torch.sigmoid(logits).cpu().numpy()
    
    results = {}
    for idx, prob in enumerate(probs):
        label = model.config.id2label[idx]
        results[label] = float(prob)
    
    return results

# Use gr.JSON output instead of gr.Label for API compatibility
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Abstract", placeholder="Enter thesis abstract...", lines=5),
    outputs=gr.JSON(label="SDG Predictions"),
    title="SDG Thesis Classifier",
    description="Multi-label classification for UN Sustainable Development Goals"
)

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