ekatra/mobius-v2

1from numpy.linalg import norm2import numpy as np3from sentence_transformers import SentenceTransformer4import PyPDF25from nltk.tokenize import sent_tokenize6 7 8def read_pdf(fname):9    """10    This function reads the pdf file and extracts the text from it.11 12    Parameters:13    fname (str): Name of the pdf file14 15    Returns:16    text_ext (list): List of extracted text from the pdf file17    """18    reader = PyPDF2.PdfReader(fname)19    text_ext = []20    for i in range(len(reader.pages)):21        pageObj = reader.pages[i]22        # extracting text from page23        text_ext.append(pageObj.extract_text())24 25    return text_ext26 27 28def sent_tokenize(text_ext):29    """30    This function apply sent_tokenize to the text and stores the result in a list.31 32    Parameters:33    text_ext (list): List of extracted text from the pdf file34 35    Returns:36    sent_toks (list): List of tokenized sentences37    """38    sent_toks = []39 40    for i in text_ext:41        sent_toks.append(sent_tokenize(i))42    print("len(sent_toks) ", len(sent_toks))43 44    return sent_toks45 46 47def create_content_embeddings(concat_list):48    """49    This function creates embeddings for the document sentences.50 51    Parameters:52    concat_list (list): List of tokenized sentences53 54    Returns:55    embeddings (list): List of embeddings of the sentences56    """57    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')58    embeddings = model.encode(concat_list)59 60    return embeddings61 62 63def create_query_embeddings(query_text):64    """65    This function creates embeddings for the query.66    Parameters:67    query_text (str): Query entered by the user68 69    Returns:70    query_embedding (list): List of embeddings of the query71    """72 73    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')74    query_embedding = model.encode(query_text)75    return query_embedding76 77 78def calculate_cosine(query_embedding, embeddings, concat_list):79    """80 81    This function calculates cosine similarity between the query and the sentences.82 83    Parameters:84    query_embedding (list): List of embeddings of the query85    embeddings (list): List of embeddings of the sentences86    concat_list (list): List of tokenized sentences87 88    Returns:89    cosine_lis (list): List of cosine similarity values90    """91    cosine_lis = []92 93    for i in range(len(concat_list)):94        cosine = np.dot(query_embedding,95                        embeddings[i]) / (norm(query_embedding)*norm(embeddings[i]))96        cosine_lis.append(cosine)97 98    # print("cosine_lis ", cosine_lis)99    return (cosine_lis)100 101 102def fetch_top_rank_ans(cosine_lis, N):103    """104    This function fetches the top N ranked sentences.105 106    Parameters:107    cosine_lis (list): List of cosine similarity values108    N (int): Number of sentences to be ranked109 110    Returns:111    indexes_final (list): List of top N ranked sentences112    """113 114    list1 = cosine_lis115    indexes_final = sorted(116        range(len(list1)), key=lambda i: list1[i], reverse=True)[:N]117 118    print("indexes_final ", indexes_final, len(indexes_final))119    indices = len(list1)120 121    sorted_indices = sorted(range(indices),122                            key=lambda i: list1[i], reverse=True)123    print(indices, indices)124 125    if indices < N:126        N = indices127    indexes_final = []128 129    for i in range(N):130        indexes_final.append(sorted_indices[i])131    len(indexes_final)132    return indexes_final133 134 135def fetch_most_relevant(indexes_final, concat_list, list1, query):136    """137    This function fetches the most relevant sentences, pass it as a context to GPT-3 prompt along with user's query.138 139    Parameters:140    indexes_final (list): List of top N ranked sentences141    concat_list (list): List of tokenized sentences142    list1 (list): List of cosine similarity values143    query (str): Query entered by the user144 145    Returns:146    prompt (str): GPT-3 prompt147    """148 149    dicts = {}150 151    keys = indexes_final152    # print(indexes_final)153    for i in keys:154        dicts[i] = concat_list[i]155 156    most_relevant_document_sections = [dicts]157 158    len(most_relevant_document_sections)159 160    chosen_sections = []161    chosen_sections_len = 0162    chosen_sections_indexes = []163 164    indices = range(len(list1))165    sorted_indices = sorted(indices, key=lambda i: list1[i], reverse=True)166    print(indexes_final, len(indexes_final))167 168    for section_index in range(len(indexes_final)):169 170        if chosen_sections_len > 500:171            break172        chosen_sections.append(173            concat_list[sorted_indices[section_index]].replace("\n", " "))174        chosen_sections_indexes.append(str(section_index))175 176    # Useful diagnostic information177    print(f"Selected {len(chosen_sections)} document sections:")178 179    header = """Answer the question as a human in natural language conversation using the provided context, and if the answer is not contained within the text below, say "I don't have that information"\n\nContext:\n"""180 181    # print(query)182    prompt = header + "".join(chosen_sections) + "\n\n Q: " + query + "\n A:"183    return prompt184