Update app.py

app.py

@@ -1,17 +1,29 @@
-import gradio as gr
 from transformers import pipeline
+import gradio as gr
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-model = 'MoritzLaurer/deberta-v3-base-zeroshot-v1.1-all-33'
-pipe = pipeline('zero-shot-classification', model=model)
+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):
-    output = pipe(text, classes, multi_label=multi_label)
-    print(output)
-    return dict(zip(output['labels'], output['scores']))
+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)
+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], outputs='label')
+app = gr.Interface(fn=infer, inputs=[text_input, class_input, allow_multi_label, model_sizes], outputs='label')
 app.launch()