File size: 1,602 Bytes
0c7add2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import faiss
import pickle
import os
from PyPDF2 import PdfReader
import glob
from pathlib import Path
import re
import requests
from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain import PromptTemplate
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
import dotenv
dotenv.load_dotenv()
class SemanticSearch():
def __init__(
self,
threshold: float,
with_source=False,
k=5,
):
self.threshold = threshold
self.with_source = with_source
self.k = k
with open('./data/store.pkl', 'rb') as f:
self.db = pickle.load(f)
def __call__(self, query):
documents = self.db.similarity_search_with_score(query, k=self.k)
if len(documents) == 0:
return None
if not self.with_source:
output = '\n\n\n'.join([i[0].page_content for i in documents])
else:
output = '\n\n\n'.join([i[0].page_content + '\n\nSource:' + os.path.basename(
str(i[0].metadata['source']) + '\n') for i in documents])
return output
class ContentSearch():
def __init__(
self,
semantic_search,
prompt_template,
):
self.semantic_search = semantic_search
self.prompt_template = prompt_template
def __call__(self, query):
content = self.semantic_search(query)
if content is None:
return "No results found"
else:
return self.prompt_template.format(content=content)
|