Spaces:
Runtime error
Runtime error
| 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", | |
| environment="us-central1-gcp" # find next to API key in console | |
| ) | |
| index_name = "abstractive-question-answering" | |
| # check if the abstractive-question-answering index exists | |
| if index_name not in pinecone.list_indexes(): | |
| # create the index if it does not exist | |
| pinecone.create_index( | |
| index_name, | |
| dimension=768, | |
| metric="cosine" | |
| ) | |
| # connect to abstractive-question-answering index we created | |
| index = pinecone.Index(index_name) | |
| from transformers import BartTokenizer, BartForConditionalGeneration | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
| model = model.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() | |
| # search pinecone index for context passage with the answer | |
| xc = index.query(xq, top_k=top_k, include_metadata=True) | |
| return xc | |
| def format_query(query, context): | |
| # extract passage_text from Pinecone search result and add the <P> tag | |
| context = [f"<P> {m['metadata']['text']}" for m in context] | |
| # concatinate all context passages | |
| context = " ".join(context) | |
| # contcatinate the query and context passages | |
| query = f"question: {query} context: {context}" | |
| return query | |
| def generate_answer(query): | |
| 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!') | |
| if b: | |
| st.write("Processing, please wait!") | |
| context = query_pinecone(query, top_k=5) | |
| query = format_query(query, context["matches"]) | |
| generate_answer(query) |