PuzzleZeroShot / app.py
belyakoff's picture
Update app.py
faf1868 verified
import os
import spaces
import gradio as gr
from transformers import pipeline
classifier = pipeline("zero-shot-classification", model=os.getenv('MODEL'))
@spaces.GPU(duration=120)
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)
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()