samsonleegh's picture
Create app.py
c078977 verified
raw
history blame
2.79 kB
import os
import gradio as gr
from dotenv import load_dotenv
from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler, CBEventType
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.llms.openai import OpenAI
from llama_index.llms.groq import Groq
from llama_index.core.base.embeddings.base import similarity
from llama_index.llms.ollama import Ollama
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core import StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import load_index_from_storage
load_dotenv()
# set up LLM
GROQ_API_KEY = os.getenv('GROQ_API_KEY')
llm = Groq(model="llama3-70b-8192")
Settings.llm = llm
# set up callback manager
llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])
Settings.callback_manager = callback_manager
# converting documents into embeddings and indexing
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.embed_model = embed_model
# create splitter
splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=20)
Settings.transformations = [splitter]
if os.path.exists("./vectordb"):
storage_context = StorageContext.from_defaults(persist_dir="./vectordb")
index = load_index_from_storage(storage_context)
else:
filename_fn = lambda filename: {"file_name": filename}
required_exts = [".pdf",".docx"]
reader = SimpleDirectoryReader(
input_dir="./data",
required_exts=required_exts,
recursive=True,
file_metadata=filename_fn
)
documents = reader.load_data()
for doc in documents:
doc.text = str(doc.metadata) +' '+ doc.text
print("index creating with `%d` documents", len(documents))
index = VectorStoreIndex.from_documents(documents, embed_model=embed_model, text_splitter=splitter)
index.storage_context.persist(persist_dir="./vectordb")
# set up query engine
query_engine = index.as_query_engine(
similarity_top_k=5,
#node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.7)],
verbose=True,
)
def retreive(question):
qns_w_source = "Answer the following question: " + question + " Followed by providing the page and file name of the source document as well, thank you!"
streaming_response = query_engine.query(qns_w_source)
#sources = streaming_response.get_formatted_sources(length=5000)
return str(streaming_response) # + "\n" + str(sources)
demo = gr.Interface(fn=retreive, inputs="textbox", outputs="textbox")
if __name__ == "__main__":
demo.launch()