import os 
import spaces
import gradio as gr
from transformers import pipeline

classifier = pipeline("zero-shot-classification", model=os.getenv('MODEL'))

@spaces.GPU(duration=100)
def classify(text, labels):
    if not text or not labels:
        return []
    labels = list(map(lambda x: x.strip(), labels.split(',')))
    result = classifier(text, candidate_labels=labels, multi_label=False)
    return list(zip(result['labels'], map(lambda x: round(x, 4), result['scores'])))

with gr.Blocks() as demo:
    gr.Markdown("## Zero-Shot классификация")
    
    with gr.Row():
        with gr.Column(scale=1.5):
            text_input = gr.Textbox(label="Текст для классификации", placeholder="Введите текст...")
            labels_input = gr.Textbox(label="Классы (через запятую)", placeholder="Опишите классы через запятую...")
            button = gr.Button("Classify")
        with gr.Column(scale=1):
            output = gr.Dataframe(headers=["Класс", "Вероятность"], label="Результаты классификации")
    
    button.click(classify, inputs=[text_input, labels_input], outputs=output)
    
demo.launch()