[UPD] softmax model

app.py

@@ -28,8 +28,8 @@ def load_model_and_tokenizer():
     
     return model, tokenizer, unique_tags
 
-def predict(texts, model, tokenizer, unique_tags, top_threshold=0.95):
-    device = next(model.parameters()).device  # Берём девайс из модели
+def predict(texts, model, tokenizer, unique_tags, top_threshold=0.95, max_tags=4):
+    device = next(model.parameters()).device
     model.eval()
     predictions = []
     
@@ -44,17 +44,18 @@ def predict(texts, model, tokenizer, unique_tags, top_threshold=0.95):
             ).to(device)
             
             outputs = model(**inputs)
-            probs = torch.sigmoid(outputs.logits).cpu().numpy().flatten()
+            raw_probs = torch.softmax(outputs.logits, dim=1).cpu().numpy().flatten()
             
-            sorted_indices = probs.argsort()[::-1]
-            cumulative = 0
+            sorted_indices = raw_probs.argsort()[::-1]
+            cumulative = 0.0
             selected = []
             
             for idx in sorted_indices:
-                cumulative += probs[idx]
-                selected.append((unique_tags[idx], float(probs[idx])))
-                if cumulative >= top_threshold:
-                    break
+                prob = raw_probs[idx]
+                cumulative += prob
+                selected.append((unique_tags[idx], float(prob)))
+                if cumulative >= top_threshold or len(selected) == max_tags:
+                    break             
             predictions.append(selected)
     
     return predictions