Update README.md

README.md

@@ -39,22 +39,21 @@ def predict_dialect(sent):
        CAMeL Tools MADAR 6 DID model"""
 
     predictions = DID.predict([sent])
+    scores = predictions[0].scores
 
     if predictions[0].top != "MSA":
-        scores = predictions[0].scores
-        highest = sorted(
-            scores.items(), key=lambda x: x[1], reverse=True)[0]
-        name = highest[0]
-        score = highest[1]
-
+        # get the highest pred
+        pred = sorted(scores.items(),
+                      key=lambda x: x[1], reverse=True)[0]
     else:
-        scores = predictions[0].scores
-        second_highest = sorted(
-            scores.items(), key=lambda x: x[1], reverse=True)[1]
-        name = second_highest[0]
-        score = second_highest[1]
+        # get the second highest pred
+        pred = sorted(scores.items(),
+                      key=lambda x: x[1], reverse=True)[1]
+
+    dialect = pred[0]
+    score = pred[1]
 
-    return name, score
+    return dialect, score
 
 tokenizer = AutoTokenizer.from_pretrained('CAMeL-Lab/arat5-coda-did')
 model = AutoModelForSeq2SeqLM.from_pretrained('CAMeL-Lab/arat5-coda-did')