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! As the user might get the answer he or she wants, you could help provide more details and relevant information referred for your answer. This will enable/prompt them to further phrase their questions." | |
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(share=True) |