Spaces:
Paused
Paused
| from sentence_transformers import SentenceTransformer, CrossEncoder, util | |
| from torch import tensor as torch_tensor | |
| from datasets import load_dataset | |
| from langchain.llms import OpenAI | |
| from langchain.docstore.document import Document | |
| from langchain.chains.question_answering import load_qa_chain | |
| from langchain.chains.qa_with_sources import load_qa_with_sources_chain | |
| from langchain.prompts import PromptTemplate | |
| """# import models""" | |
| bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') | |
| bi_encoder.max_seq_length = 256 #Truncate long passages to 256 tokens | |
| #The bi-encoder will retrieve top_k documents. We use a cross-encoder, to re-rank the results list to improve the quality | |
| cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') | |
| """# import datasets""" | |
| dataset = load_dataset("gfhayworth/hack_policy", split='train') | |
| mypassages = list(dataset.to_pandas()['psg']) | |
| dataset_embed = load_dataset("gfhayworth/hack_policy_embed", split='train') | |
| dataset_embed_pd = dataset_embed.to_pandas() | |
| mycorpus_embeddings = torch_tensor(dataset_embed_pd.values) | |
| def search(query, passages = mypassages, doc_embedding = mycorpus_embeddings, top_k=20, top_n = 1): | |
| question_embedding = bi_encoder.encode(query, convert_to_tensor=True) | |
| question_embedding = question_embedding #.cuda() | |
| hits = util.semantic_search(question_embedding, doc_embedding, top_k=top_k) | |
| hits = hits[0] # Get the hits for the first query | |
| ##### Re-Ranking ##### | |
| cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits] | |
| cross_scores = cross_encoder.predict(cross_inp) | |
| # Sort results by the cross-encoder scores | |
| for idx in range(len(cross_scores)): | |
| hits[idx]['cross-score'] = cross_scores[idx] | |
| hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) | |
| predictions = hits[:top_n] | |
| return predictions | |
| # for hit in hits[0:3]: | |
| # print("\t{:.3f}\t{}".format(hit['cross-score'], mypassages[hit['corpus_id']].replace("\n", " "))) | |
| def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings): | |
| predictions = search(qry, passages = passages, doc_embedding = doc_embedding, top_n=5, ) | |
| prediction_text = [] | |
| for hit in predictions: | |
| page_content = passages[hit['corpus_id']] | |
| metadata = {"source": hit['corpus_id']} | |
| result = Document(page_content=page_content, metadata=metadata) | |
| prediction_text.append(result) | |
| return prediction_text | |
| template = """You are a friendly AI assistant for the insurance company Humana. Given the following extracted parts of a long document and a question, create a succinct final answer. | |
| If you don't know the answer, just say that you don't know. Don't try to make up an answer. | |
| If the question is not about Humana, politely inform them that you are tuned to only answer questions about Humana. | |
| QUESTION: {question} | |
| ========= | |
| {context} | |
| ========= | |
| FINAL ANSWER:""" | |
| PROMPT = PromptTemplate(template=template, input_variables=["context", "question"]) | |
| chain_qa = load_qa_chain(OpenAI(temperature=0), chain_type="stuff", prompt=PROMPT) | |
| def get_llm_response(message): | |
| mydocs = get_text_fmt(message) | |
| response = chain_qa.run(input_documents=mydocs, question=message) | |
| return response | |
| def chat(message, history): | |
| history = history or [] | |
| message = message.lower() | |
| response = get_llm_response(message) | |
| history.append((message, response)) | |
| return history, history | |