Update app.py

app.py

@@ -4,6 +4,7 @@ from sentence_transformers import SentenceTransformer
 import pickle
 import re
 from transformers import pipeline
+import torch
 
 
 
@@ -22,19 +23,21 @@ with open('articles.pkl', 'rb') as file:
 
 index_loaded = faiss.read_index("sentence_embeddings_index_no_citation.faiss")
 
+
+device = 0 if torch.cuda.is_available() else -1
 if 'model_embedding' not in st.session_state:
-    st.session_state.model_embedding = SentenceTransformer('bkai-foundation-models/vietnamese-bi-encoder')
+    st.session_state.model_embedding = SentenceTransformer('bkai-foundation-models/vietnamese-bi-encoder', device = device)
 
 
 
 # Replace this with your own checkpoint
 model_checkpoint = "model"
-question_answerer = pipeline("question-answering", model=model_checkpoint)
+question_answerer = pipeline("question-answering", model=model_checkpoint, device = device)
 def question_answering(question):
     print(question)
     query_sentence = [question]
     query_embedding = st.session_state.model_embedding.encode(query_sentence)
-    k = 10
+    k = 20
     D, I = index_loaded.search(query_embedding.astype('float32'), k)  # D is distances, I is indices
     answer = [question_answerer(question=query_sentence[0], context=articles[I[0][i]], max_answer_len = 256) for i in range(k)]
     best_answer = max(answer, key=lambda x: x['score'])