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

model_name = "nattkorat/xlm-roberta-with-raw-text-for-text-classification"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = pipeline(task="text-classification", 
                 model=model_name, 
                 tokenizer=tokenizer, 
                 truncation=True, 
                 padding=True, 
                 top_k=6
                 )



def classify_news(article: str):
    if not article.strip():
        return {"error": "សូមបញ្ចូលអត្ថបទដើម្បីវិភាគ!"}
    output = model(article)
    results = {}
    for out in output[0]:
        results[out['label']] = out['score']
    return results

interface = gr.Interface(
    fn=classify_news,
    inputs=gr.TextArea(placeholder="ដាក់អត្ថបទចូលនៅទីនេះ", label="អត្ថបទ"),
    outputs=gr.Label(num_top_classes=6),
    title="XLM-Roberta for Khmer News Classification",
    description="A demo of XLM-Roberta for khmer news classification."
)

interface.launch()


# print(classify_news("Hello"))