import gradio as gr
from transformers import AutoModel, AutoTokenizer
import json

# Load model and tokenizer
repo_id = "rasoultilburg/SocioCausaNet"
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(repo_id)

# Prediction function
def predict(sentences, rel_mode="auto", rel_threshold=0.5, cause_decision="cls+span"):
    results = model.predict(
        sentences,
        tokenizer=tokenizer,
        rel_mode=rel_mode,
        rel_threshold=rel_threshold,
        cause_decision=cause_decision
    )
    return json.dumps(results, indent=2, ensure_ascii=False)

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Sentences (comma-separated)", placeholder="Enter sentences"),
        gr.Radio(["auto", "neural_only"], label="Relation Mode", value="auto"),
        gr.Slider(0.0, 1.0, value=0.5, label="Relation Threshold"),
        gr.Radio(["cls_only", "span_only", "cls+span"], label="Cause Decision", value="cls+span")
    ],
    outputs="text",
    title="SocioCausaNet API",
    description="Extract causal relations from text"
)

iface.launch()