from transformers import pipeline
import gradio as gr
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

pipelines = {
    'small': pipeline('zero-shot-classification', model='MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33'),
    'base': pipeline('zero-shot-classification', model='MoritzLaurer/deberta-v3-base-zeroshot-v1.1-all-33'),
    'large': pipeline('zero-shot-classification', model='MoritzLaurer/deberta-v3-large-zeroshot-v1.1-all-33')
}

def infer(text, classes, multi_label, model_size):
    try:
        output = pipelines[model_size](text, classes, multi_label=multi_label)
        logger.info(f"Model size: {model_size}, Output: {output}")
        return dict(zip(output['labels'], output['scores']))
    except Exception as e:
        logger.error(f"Error: {e}")
        return {}

text_input = gr.Textbox(lines=5, placeholder='Once upon a time...', label='Text Source', show_label=True)
class_input = gr.Textbox(value='positive, negative', label='Class Label', show_label=True, info='Use commas (,) to seperate classes')
allow_multi_label = gr.Checkbox(value=True, label='Multiple True Classes')
model_sizes = gr.Radio(choices=['small', 'base', 'large'], value='base', label='Model Sizes', show_label=True)

app = gr.Interface(fn=infer, inputs=[text_input, class_input, allow_multi_label, model_sizes], outputs='label')
app.launch()