Update app.py

app.py

@@ -311,6 +311,7 @@ def chat_bert_context(question, history):
     
 #-------------------------------------Bi-BERT-Encoder------------------------------------------#
 MAX_LENGTH = 128
+inverted_answer = dict(enumerate(df.answer.tolist()))
 # Define function for mean-pooling
 def mean_pool(token_embeds: torch.tensor, attention_mask: torch.tensor) -> torch.tensor:
     in_mask = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()
@@ -372,12 +373,13 @@ def chat_bi_bert(question):
     cosine_similarities = cosine_similarity([question], question_embeds).flatten()
     top_indice = np.argmax(cosine_similarities, axis=0)
     answer = df['answer'].iloc[top_indice]
+    answer = inverted_answer[top_indice]
     return answer
 
 
 
 #-------------------------------------Bi+Cross-BERT-Encoder------------------------------------------#
-inverted_answer = dict(enumerate(df.answer.tolist()))
+
 
 #Define class for CrossEncoderBert
 class CrossEncoderBert(torch.nn.Module):
@@ -418,9 +420,9 @@ def chat_cross_bert(question):
 
     # Process scores for finetuned model
     scores = ce_scores.cpu().numpy()
-    scores_ix = np.argmax(scores)
+    ix = np.argmax(scores)
     # print(f"{corpus[scores_ix]}")
-    return corpus[scores_ix]
+    return corpus[ix]
 
 # gradio part 
 def echo(message, history, model):