conbein/lc-server

1import os2import re3import shutil4import urllib.request5from pathlib import Path6from tempfile import NamedTemporaryFile7 8import fitz9import numpy as np10import openai11import tensorflow_hub as hub12from fastapi import UploadFile13from lcserve import serving14from sklearn.neighbors import NearestNeighbors15 16 17recommender = None18 19 20def download_pdf(url, output_path):21    urllib.request.urlretrieve(url, output_path)22 23 24def preprocess(text):25    text = text.replace('\n', ' ')26    text = re.sub('\s+', ' ', text)27    return text28 29 30def pdf_to_text(path, start_page=1, end_page=None):31    doc = fitz.open(path)32    total_pages = doc.page_count33 34    if end_page is None:35        end_page = total_pages36 37    text_list = []38 39    for i in range(start_page - 1, end_page):40        text = doc.load_page(i).get_text("text")41        text = preprocess(text)42        text_list.append(text)43 44    doc.close()45    return text_list46 47 48def text_to_chunks(texts, word_length=150, start_page=1):49    text_toks = [t.split(' ') for t in texts]50    page_nums = []51    chunks = []52 53    for idx, words in enumerate(text_toks):54        for i in range(0, len(words), word_length):55            chunk = words[i : i + word_length]56            if (57                (i + word_length) > len(words)58                and (len(chunk) < word_length)59                and (len(text_toks) != (idx + 1))60            ):61                text_toks[idx + 1] = chunk + text_toks[idx + 1]62                continue63            chunk = ' '.join(chunk).strip()64            chunk = f'[Page no. {idx+start_page}]' + ' ' + '"' + chunk + '"'65            chunks.append(chunk)66    return chunks67 68 69class SemanticSearch:70    def __init__(self):71        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')72        self.fitted = False73 74    def fit(self, data, batch=1000, n_neighbors=5):75        self.data = data76        self.embeddings = self.get_text_embedding(data, batch=batch)77        n_neighbors = min(n_neighbors, len(self.embeddings))78        self.nn = NearestNeighbors(n_neighbors=n_neighbors)79        self.nn.fit(self.embeddings)80        self.fitted = True81 82    def __call__(self, text, return_data=True):83        inp_emb = self.use([text])84        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]85 86        if return_data:87            return [self.data[i] for i in neighbors]88        else:89            return neighbors90 91    def get_text_embedding(self, texts, batch=1000):92        embeddings = []93        for i in range(0, len(texts), batch):94            text_batch = texts[i : (i + batch)]95            emb_batch = self.use(text_batch)96            embeddings.append(emb_batch)97        embeddings = np.vstack(embeddings)98        return embeddings99 100 101def load_recommender(path, start_page=1):102    global recommender103    if recommender is None:104        recommender = SemanticSearch()105 106    texts = pdf_to_text(path, start_page=start_page)107    chunks = text_to_chunks(texts, start_page=start_page)108    recommender.fit(chunks)109    return 'Corpus Loaded.'110 111 112def generate_text(openAI_key, prompt, engine="text-davinci-003"):113    openai.api_key = openAI_key114    completions = openai.Completion.create(115        engine=engine,116        prompt=prompt,117        max_tokens=512,118        n=1,119        stop=None,120        temperature=0.7,121    )122    message = completions.choices[0].text123    return message124 125 126def generate_answer(question, openAI_key):127    topn_chunks = recommender(question)128    prompt = ""129    prompt += 'search results:\n\n'130    for c in topn_chunks:131        prompt += c + '\n\n'132 133    prompt += (134        "Instructions: Compose a comprehensive reply to the query using the search results given. "135        "Cite each reference using [ Page Number] notation (every result has this number at the beginning). "136        "Citation should be done at the end of each sentence. If the search results mention multiple subjects "137        "with the same name, create separate answers for each. Only include information found in the results and "138        "don't add any additional information. Make sure the answer is correct and don't output false content. "139        "If the text does not relate to the query, simply state 'Text Not Found in PDF'. Ignore outlier "140        "search results which has nothing to do with the question. Only answer what is asked. The "141        "answer should be short and concise. Answer step-by-step. \n\nQuery: {question}\nAnswer: "142    )143 144    prompt += f"Query: {question}\nAnswer:"145    answer = generate_text(openAI_key, prompt, "text-davinci-003")146    return answer147 148 149def load_openai_key() -> str:150    key = os.environ.get("OPENAI_API_KEY")151    if key is None:152        raise ValueError(153            "[ERROR]: Please pass your OPENAI_API_KEY. Get your key here : https://platform.openai.com/account/api-keys"154        )155    return key156 157 158@serving159def ask_url(url: str, question: str):160    download_pdf(url, 'corpus.pdf')161    load_recommender('corpus.pdf')162    openAI_key = load_openai_key()163    return generate_answer(question, openAI_key)164 165 166@serving167async def ask_file(file: UploadFile, question: str) -> str:168    suffix = Path(file.filename).suffix169    with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:170        shutil.copyfileobj(file.file, tmp)171        tmp_path = Path(tmp.name)172 173    load_recommender(str(tmp_path))174    openAI_key = load_openai_key()175    return generate_answer(question, openAI_key)176