Update app.py

app.py

@@ -27,6 +27,14 @@ from transformers import BartTokenizer, BartForConditionalGeneration
 tokenizer = BartTokenizer.from_pretrained('vblagoje/bart_lfqa')
 generator = BartForConditionalGeneration.from_pretrained('vblagoje/bart_lfqa').to('cpu')
 
+import torch
+from sentence_transformers import SentenceTransformer
+
+# set device to GPU if available
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# load the retriever model from huggingface model hub
+retriever = SentenceTransformer("flax-sentence-embeddings/all_datasets_v3_mpnet-base", device=device)
+
 def query_pinecone(query, top_k):
     # generate embeddings for the query
     xq = retriever.encode([query]).tolist()