import gradio as gr
from transformers import pipeline

classifier = pipeline(
    "text-generation",
    model="AtrriJi/smolified-risk-clause-classifier",
    device=-1
)

def classify(clause):

    prompt = f"""
Classify the legal clause into one of:
Payment Terms, Intellectual Property, Confidentiality,
Termination, Indemnification, Force Majeure,
Governing Law, Warranty, Limitation of Liability,
Dispute Resolution.

Also predict risk level: Low, Medium, or High.

Return ONLY JSON like:
{{"category": "...", "risk_level": "..."}}

Clause:
{clause}
"""

    result = classifier(
        prompt,
        max_new_tokens=100,
        do_sample=False
    )[0]["generated_text"]

    # Remove prompt part
    clean_output = result[len(prompt):].strip()

    return clean_output


demo = gr.Interface(
    fn=classify,
    inputs="text",
    outputs="text",
    title="Risk Clause Classifier"
)

demo.launch()