Update bartpho/utils.py

bartpho/utils.py

@@ -51,28 +51,27 @@ def predict(model, text, tokenizer, model_tokenize=None, processed=True, printou
     for i in range(no_tag):
         tag = tags[i]
         score_list = []
-        input_ids = tokenizer([text] * no_polarity, return_tensors='pt')['input_ids']
+
+        input_ids = tokenizer([text] * no_polarity, return_tensors='pt')['input_ids'].to(device)
         target_list = ["Nhận_xét " + tag.lower() + " " + polarity.lower() + " ." for polarity in polarity_list]
-        output_ids = tokenizer(target_list, return_tensors='pt', padding=True, truncation=True)['input_ids']
+        output_ids = tokenizer(target_list, return_tensors='pt', padding=True, truncation=True)['input_ids'].to(device)
 
         with torch.no_grad():
             output = model(input_ids=input_ids.to(device), decoder_input_ids=output_ids.to(device))[0]
             logits = output.softmax(dim=-1).to('cpu').numpy()
         for m in range(no_polarity):
-            score = 1
-            for n in range(logits[m].shape[0] - 2):
-                score *= logits[m][n][output_ids[m][n+1]]
+            score = np.sum(np.log(logits[m][range(len(output_ids[m]) - 2), output_ids[m][1:-1]]))
             score_list.append(score)
-        predict = np.argmax(score_list)
+        predict = int(np.argmax(score_list))  # Ép kiểu sang int
         predicts.append(predict)
         
     if printout:
         result = {}
         for i in range(no_tag):
-            if predicts[i] != 0:
-                result[eng_tags[i]] = eng_polarity[predicts[i]]
-        print(result)
-    return predicts
+            if predicts[i] != 0:  # Bỏ qua các nhãn không có cảm xúc (mặc định 0)
+                result[tags[i]] = polarity_list[predicts[i]]  # Ánh xạ nhãn
+        # print(result)
+    return result
 
 def predict_df(model, df, tokenizer=None, model_tokenize=None, tokenizer_name='vinai/bartpho-word-base', processed=True, printout=True):
     model.eval()