import gradio as gr
from transformers import pipeline
import spaces # 1. Import the spaces library

model_id = "AndreiUrsu/MediaManipulationDetection" 
classifier = pipeline("text-classification", model=model_id)

@spaces.GPU # 2. Add the decorator to request GPU allocation
def clasifica_stire(headline, continut):
    text_complet = f"{headline} {continut}"
    rezultat = classifier(text_complet)[0]
    
    eticheta_bruta = str(rezultat['label'])
    scor = rezultat['score']
    
    # Maparea etichetelor brute la textul dorit
    mapare_etichete = {
        "LABEL_0": "Știre Reală",
        "LABEL_1": "Știre Falsă",
        "0": "Știre Reală",
        "1": "Știre Falsă"
    }
    
    # Extragem traducerea, sau păstrăm eticheta originală dacă nu se potrivește cu maparea
    eticheta_tradusa = mapare_etichete.get(eticheta_bruta, eticheta_bruta)
    
    return f"Rezultat: {eticheta_tradusa} (Încredere: {scor:.1%})"

demo = gr.Interface(
    fn=clasifica_stire,
    inputs=[
        gr.Textbox(label="Headline / Titlu Știre", lines=1),
        gr.Textbox(label="Conținut / Corp Text", lines=5)
    ],
    outputs=gr.Textbox(label="Rezultat Evaluare"),
    title="Detector de Manipulare Media",
    description="Analizator de text în limba română pentru detectarea manipulării în știri."
)

demo.launch()