using new model now

app.py

@@ -14,18 +14,24 @@ USED_MODEL = "distilbert-base-cased"
 def load_model():
     # csv локально прочитать очень быстро, так что его не кешируем, хотя это не сложно было бы добавить наверное
     arxiv_topics_df = pd.read_csv('arxiv_topics.csv')
-    tag_to_index = {}
+    category_to_index = {}
+    current_index = 0
     for i, row in arxiv_topics_df.iterrows():
-        tag_to_index[row['tag']] = i
-    index_to_tag = {value: key for key, value in tag_to_index.items()}
-
-    return AutoModelForSequenceClassification.from_pretrained(
-        "bumchik2/train_distilbert-base-cased-tags-classification-simple", 
+        category = row['category']
+        if category not in category_to_index:
+            category_to_index[category] = current_index
+            current_index += 1
+    index_to_category = {value: key for key, value in category_to_index.items()}
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bumchik2/train-distilbert-base-cased-tags-classification", 
         problem_type="multi_label_classification", 
-        num_labels=len(tag_to_index),
-        id2label=index_to_tag,
-        label2id=tag_to_index
+        num_labels=len(category_to_index),
+        id2label=index_to_category,
+        label2id=category_to_index
     )
+    model.eval()
+    return model
 
 model = load_model()
 
@@ -44,21 +50,23 @@ def tokenize_function(text):
 def get_category_probs_dict(model, title: str, summary: str) -> Dict[str, float]:
     # csv локально прочитать очень быстро, так что его не кешируем, хотя это не сложно было бы добавить наверное
     arxiv_topics_df = pd.read_csv('arxiv_topics.csv')
-    tag_to_index = {}
-    tag_to_category = {}
+    category_to_index = {}
+    current_index = 0
     for i, row in arxiv_topics_df.iterrows():
-        tag_to_category[row['tag']] = row['category']
-        tag_to_index[row['tag']] = i
-    index_to_tag = {value: key for key, value in tag_to_index.items()}
+        category = row['category']
+        if category not in category_to_index:
+            category_to_index[category] = current_index
+            current_index += 1
+    index_to_category = {value: key for key, value in category_to_index.items()}
 
     text = f'{title} $ {summary}'
-    tags_logits = model(**{key: torch.tensor(value).to(model.device).unsqueeze(0) for key, value in tokenize_function(text).items()}).logits
+    category_logits = model(**{key: torch.tensor(value).to(model.device).unsqueeze(0) for key, value in tokenize_function(text).items()}).logits
     sigmoid = torch.nn.Sigmoid()
-    tags_probs = sigmoid(tags_logits.squeeze().cpu()).numpy()
-    tags_probs /= tags_probs.sum()
+    category_probs = sigmoid(category_logits.squeeze().cpu()).numpy()
+    category_probs /= category_probs.sum()
     category_probs_dict = {category: 0.0 for category in set(arxiv_topics_df['category'])}
-    for index in range(len(index_to_tag)):
-        category_probs_dict[tag_to_category[index_to_tag[index]]] += float(tags_probs[index])
+    for index in range(len(index_to_category)):
+        category_probs_dict[index_to_category[index]] += float(category_probs[index])
     return category_probs_dict