Update app.py

app.py

@@ -71,6 +71,9 @@ bert_tokenizer = AutoTokenizer.from_pretrained("my_finetuned_model")
 bert_model = AutoModelForSequenceClassification.from_pretrained("my_finetuned_model")
 bert_model.eval()
 
+# --- Rječnik za mapiranje oznaka ---
+label_names = {0: 'pozitivno', 1: 'neutralno', 2: 'negativno'}
+
 # --- Pretvaranje teksta u indekse za CNN i GRU ---
 def text_to_indices(text, max_len=100):
     tokens = text.lower().split()
@@ -92,7 +95,7 @@ def predict_svm(text):
     proba = svm_pipeline.predict_proba([text])[0]
     pred = svm_pipeline.classes_[proba.argmax()]
     print(f"SVM predikcija: {pred}, povjerenje: {proba.max():.2f}")
-    return f"{pred} (p={proba.max():.2f})"
+    return f"{label_names[pred]} (p={proba.max():.2f})"
 
 def predict_cnn(text):
     print(f"Predikcija CNN za tekst: {text}")
@@ -104,7 +107,7 @@ def predict_cnn(text):
         pred = torch.argmax(probs, dim=1).item()
         confidence = probs[0][pred].item()
     print(f"CNN predikcija: {pred}, povjerenje: {confidence:.2f}")
-    return f"{pred} (p={confidence:.2f})"
+    return f"{label_names[pred]} (p={confidence:.2f})"
 
 def predict_gru(text):
     print(f"Predikcija GRU za tekst: {text}")
@@ -116,7 +119,7 @@ def predict_gru(text):
         pred = torch.argmax(probs, dim=1).item()
         confidence = probs[0][pred].item()
     print(f"GRU predikcija: {pred}, povjerenje: {confidence:.2f}")
-    return f"{pred} (p={confidence:.2f})"
+    return f"{label_names[pred]} (p={confidence:.2f})"
 
 def predict_bert(text):
     print(f"Predikcija BERTić za tekst: {text}")
@@ -128,7 +131,7 @@ def predict_bert(text):
         pred = torch.argmax(probs, dim=1).item()
         confidence = probs[0][pred].item()
     print(f"BERTić predikcija: {pred}, povjerenje: {confidence:.2f}")
-    return f"{pred} (p={confidence:.2f})"
+    return f"{label_names[pred]} (p={confidence:.2f})"
 
 # --- Gradio sučelje ---
 def predict_all(text):
@@ -149,10 +152,7 @@ demo = gr.Interface(
         gr.Textbox(label="BERTić")
     ],
     title="Demo klasifikacije teksta",
-    description=(
-        "Predikcije koriste SVM, CNN, GRU i BERTić modele.\n\n"
-        "Napomena: 0 = pozitivno, 1 = neutralno, 2 = negativno."
-    )
+    description="Predikcije koriste SVM, CNN, GRU i BERTić modele."
 )
 
 if __name__ == "__main__":