Spaces:
Sleeping
Sleeping
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() |