greg_funcs.py

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.qa_with_sources import load_qa_with_sources_chain





"""# 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


chain_qa = load_qa_with_sources_chain(OpenAI(temperature=0), chain_type="stuff")


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