Update app.py

app.py

@@ -1,6 +1,9 @@
 import pinecone
 from pprint import pprint
 import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM
+model_name = "vblagoje/bart_lfqa"
 # connect to pinecone environment
 pinecone.init(
     api_key="e5d4972e-0045-43d5-a55e-efdeafe442dd",
@@ -23,9 +26,9 @@ index = pinecone.Index(index_name)
 
 from transformers import BartTokenizer, BartForConditionalGeneration
 
-# load bart tokenizer and model from huggingface
-tokenizer = BartTokenizer.from_pretrained('vblagoje/bart_lfqa')
-generator = BartForConditionalGeneration.from_pretrained('vblagoje/bart_lfqa').to('cpu')
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+model = model.to('cuda')
 
 import torch
 from sentence_transformers import SentenceTransformer
@@ -51,13 +54,25 @@ def format_query(query, context):
     query = f"question: {query} context: {context}"
     return query
 def generate_answer(query):
-    # tokenize the query to get input_ids
-    inputs = tokenizer([query], trunication=True, max_length=1024, return_tensors="pt")
-    # use generator to predict output ids
-    ids = generator.generate(inputs["input_ids"], num_beams=2, min_length=20, max_length=64)
-    # use tokenizer to decode the output ids
-    answer = tokenizer.batch_decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    st.write(str(answer))
+    query_and_docs = query
+
+    model_input = tokenizer(query_and_docs, truncation=True, padding=True, return_tensors="pt")
+
+    generated_answers_encoded = model.generate(input_ids=model_input["input_ids"].to(device),
+                                           attention_mask=model_input["attention_mask"].to(device),
+                                           min_length=64,
+                                           max_length=256,
+                                           do_sample=False, 
+                                           early_stopping=True,
+                                           num_beams=8,
+                                           temperature=1.0,
+                                           top_k=None,
+                                           top_p=None,
+                                           eos_token_id=tokenizer.eos_token_id,
+                                           no_repeat_ngram_size=3,
+                                           num_return_sequences=1)
+    res = tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True,clean_up_tokenization_spaces=True)
+    st.write(str(res))
 
 query = st.text_area('Enter Question:')
 b = st.button('Submit!')