freeze version

app.py

@@ -1,13 +1,19 @@
 from html import unescape
 from unicodedata import normalize
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoModel
 import re
 
 re_multispace = re.compile(r"\s+")
+model_task_mapping = {
+    "Server": "Server",
+    "Category": "Category",
+    "Gender": "Gender",
+    "Day Of Week": "Day_of_week"
+}
 
 def normalize_text(text):
-    if text == None:
+    if text is None:
         return None
 
     text = text.strip()
@@ -20,28 +26,24 @@ def normalize_text(text):
     return text
 
 
-models = [
-    "Server", "Category", "Gender", "Day Of Week"
-]
-
-pipelines = {model: pipeline(task="text-classification",
- model=f"hynky/{model.replace(' ', '_')}", tokenizer="ufal/robeczech-base",
+pipelines = {task: pipeline(task="text-classification",
+ model=f"hynky/{model}", tokenizer="ufal/robeczech-base",
  truncation=True, max_length=512,
  top_k=5
-) for model in models}
+) for task, model in model_task_mapping.items()}
 
 
 def predict(article):
     article = normalize_text(article)
-    predictions = [pipelines[model](article)[0] for model in models]
+    predictions = [pipelines[model](article)[0] for model in model_task_mapping.keys()]
     predictions = [{pred["label"]: round(pred["score"], 3) for pred in task_preds} for task_preds in predictions]
-    return tuple(predictions)
+    return predictions
 
 gr.Interface(
     predict,
-    inputs=gr.inputs.Textbox(lines=4, placeholder="Paste a news article here..."),
+    inputs=gr.Textbox(lines=4, placeholder="Paste a news article here..."),
     # multioutput of gradio text
-    outputs=[gr.outputs.Label(num_top_classes=5, label=model)
-    for model in models],
+    outputs=[gr.Label(num_top_classes=5, label=task)
+    for task in model_task_mapping.keys()],
     title="News Article Classifier",
 ).launch()

requirements.txt

@@ -1,3 +1,3 @@
-transformers
-torch
-gradio==3.26.0
+transformers==0.1.1
+torch==2.1.0
+gradio==0.3.26