Update README.md

README.md

@@ -26,6 +26,7 @@ from src.modeling_enc_t5 import EncT5ForSequenceClassification
 from src.tokenization_enc_t5 import EncT5Tokenizer
 import torch
 import torch.nn.functional as F
+import numpy as np
 
 # load TART full and tokenizer
 model = EncT5ForSequenceClassification.from_pretrained("facebook/tart-full-flan-t5-xl")
@@ -49,10 +50,10 @@ print([p_1, p_2][np.argmax(normalized_scores)]) # "The population of Japan's cap
 # 2. TART-full can identify the document that is more relevant AND follows instructions.
 in_sim = "You need to find duplicated questions in Wiki forum. Could you find a question that is similar to this question"
 q_1 = "How many people live in Tokyo?"
-features = tokenizer(['{0} [SEP] {1}'.format(in_sim, q), '{0} [SEP] {1}'.format(in_sim, q)], [p, q_1], padding=True, truncation=True, return_tensors="pt")
+features = tokenizer(['{0} [SEP] {1}'.format(in_sim, q), '{0} [SEP] {1}'.format(in_sim, q)], [p_1, q_1], padding=True, truncation=True, return_tensors="pt")
 with torch.no_grad():
     scores = model(**features).logits
     normalized_scores = [float(score[1]) for score in F.softmax(scores, dim=1)]
 
-print([p, q_1][np.argmax(normalized_scores)]) #  "How many people live in Tokyo?"
+print([p_1, q_1][np.argmax(normalized_scores)]) #  "How many people live in Tokyo?"
 ```