Added locally trained models

app.py

@@ -38,11 +38,19 @@ classification_tokenizer_normal = AutoTokenizer.from_pretrained(
 classification_model_normal = AutoModelForSequenceClassification.from_pretrained(
     "sileod/deberta-v3-base-tasksource-nli"
 )
-classification_tokenizer_onnx = AutoTokenizer.from_pretrained(
-    "optimum/distilbert-base-uncased-mnli"
+
+audience_classification_tokenizer = AutoTokenizer.from_pretrained(
+    "bertugmirasyedi/deberta-v3-base-book-classification"
 )
-classification_model_onnx = ORTModelForSequenceClassification.from_pretrained(
-    "optimum/distilbert-base-uncased-mnli"
+audience_classification_model = AutoModelForSequenceClassification.from_pretrained(
+    "bertugmirasyedi/deberta-v3-base-level-classification"
+)
+
+level_classification_tokenizer = AutoTokenizer.from_pretrained(
+    "bertugmirasyedi/deberta-v3-base-level-classification"
+)
+level_classification_model = AutoModelForSequenceClassification.from_pretrained(
+    "bertugmirasyedi/deberta-v3-base-level-classification"
 )
 
 # Define similarity model
@@ -365,37 +373,68 @@ async def classify(data: list, runtime: str = "normal"):
         # Define the zero-shot classifier
         tokenizer = classification_tokenizer_normal
         model = classification_model_normal
-    elif runtime == "onnxruntime":
-        tokenizer = classification_tokenizer_onnx
-        model = classification_model_onnx
-
-    classifier_pipe = pipeline(
-        "zero-shot-classification",
-        model=model,
-        tokenizer=tokenizer,
-        hypothesis_template="This book is {}.",
-        batch_size=1,
-        device=-1,
-        multi_label=False,
-    )
 
-    # Define the candidate labels
-    level = [
-        "Introductory",
-        "Advanced",
-    ]
+        classifier_pipe = pipeline(
+            "zero-shot-classification",
+            model=model,
+            tokenizer=tokenizer,
+            hypothesis_template="This book is {}.",
+            batch_size=1,
+            device=-1,
+            multi_label=False,
+        )
 
-    audience = ["Academic", "Not Academic", "Manual"]
+        # Define the candidate labels
+        level = [
+            "Introductory",
+            "Advanced",
+        ]
 
-    classes = [
-        {
-            "audience": classifier_pipe(doc, audience)["labels"][0],
-            "audience_confidence": classifier_pipe(doc, audience)["scores"][0],
-            "level": classifier_pipe(doc, level)["labels"][0],
-            "level_confidence": classifier_pipe(doc, level)["scores"][0],
-        }
-        for doc in combined_data
-    ]
+        audience = ["Academic", "Not Academic", "Manual"]
+
+        classes = [
+            {
+                "audience": classifier_pipe(doc, audience)["labels"][0],
+                "audience_confidence": classifier_pipe(doc, audience)["scores"][0],
+                "level": classifier_pipe(doc, level)["labels"][0],
+                "level_confidence": classifier_pipe(doc, level)["scores"][0],
+            }
+            for doc in combined_data
+        ]
+    elif runtime == "local":
+        ### Define the classifier for audience prediction ###
+        audience_classifier = pipeline(
+            "text-classification",
+            model=audience_classification_model,
+            tokenizer=audience_classification_tokenizer,
+            device=-1,
+        )
+
+        ### Define the classifier for level prediction ###
+        level_classifier = pipeline(
+            "text-classification",
+            model=level_classification_model,
+            tokenizer=level_classification_tokenizer,
+            device=-1,
+        )
+
+        classes = [
+            {
+                "audience": audience_classifier(doc, padding=True, truncation=True)[0][
+                    "label"
+                ],
+                "audience_confidence": audience_classifier(
+                    doc, padding=True, truncation=True
+                )[0]["score"],
+                "level": level_classifier(doc, padding=True, truncation=True)[0][
+                    "label"
+                ],
+                "level_confidence": level_classifier(
+                    doc, padding=True, truncation=True
+                )[0]["score"],
+            }
+            for doc in combined_data
+        ]
 
     return classes