87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
import gradio as gr
|
|
|
|
import sys
|
|
import os
|
|
from collections.abc import Iterable
|
|
|
|
from langchain.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader, UnstructuredHTMLLoader
|
|
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
|
|
|
|
from langchain.chains import RetrievalQA
|
|
|
|
from langchain.llms import Ollama
|
|
|
|
from langchain.vectorstores import FAISS, Chroma
|
|
|
|
from langchain.embeddings import GPT4AllEmbeddings, CacheBackedEmbeddings
|
|
|
|
from langchain.storage import LocalFileStore#, RedisStore, UpstashRedisStore, InMemoryStore
|
|
|
|
ollama = Ollama(base_url='http://localhost:11434',
|
|
#model="codellama")
|
|
#model="starcoder")
|
|
model="llama2")
|
|
|
|
docsUrl = "/home/user/dev/docs"
|
|
|
|
documents = []
|
|
for file in os.listdir(docsUrl):
|
|
|
|
if file.endswith(".pdf"):
|
|
pdf_path = docsUrl + "/" + file
|
|
loader = PyPDFLoader(pdf_path)
|
|
documents.extend(loader.load())
|
|
print("Found " + pdf_path)
|
|
elif file.endswith('.docx') or file.endswith('.doc'):
|
|
doc_path = docsUrl + "/" + file
|
|
loader = Docx2txtLoader(doc_path)
|
|
documents.extend(loader.load())
|
|
print("Found " + doc_path)
|
|
elif file.endswith('.txt') or file.endswith('.kt') or file.endswith('.json'):
|
|
text_path = docsUrl + "/" + file
|
|
loader = TextLoader(text_path)
|
|
documents.extend(loader.load())
|
|
print("Found " + text_path)
|
|
elif file.endswith('.html') or file.endswith('.htm'):
|
|
htm_path = docsUrl + "/" + file
|
|
loader = UnstructuredHTMLLoader(htm_path)
|
|
documents.extend(loader.load())
|
|
print("Found " + htm_path)
|
|
|
|
|
|
text_splitter = CharacterTextSplitter(chunk_size=32, chunk_overlap=32)
|
|
all_splits = text_splitter.split_documents(documents)
|
|
|
|
|
|
|
|
#fs = LocalFileStore("/home/gabriele/dev/cache/")
|
|
|
|
#underlying_embeddings = GPT4AllEmbeddings()
|
|
#cached_embedder = CacheBackedEmbeddings.from_bytes_store(
|
|
# underlying_embeddings, fs, namespace=underlying_embeddings.model
|
|
#)
|
|
|
|
|
|
|
|
vectorstore = Chroma.from_documents(documents=all_splits, embedding=GPT4AllEmbeddings(embeddings_chunk_size=1000))
|
|
#vectorstore = FAISS.from_documents(documents=all_splits, embedding=cached_embedder)
|
|
|
|
|
|
def AI_response(question, history):
|
|
docs = vectorstore.similarity_search(question)
|
|
len(docs)
|
|
qachain=RetrievalQA.from_chain_type(ollama, retriever=vectorstore.as_retriever())
|
|
#reply=qachain()
|
|
#reply=str(qachain({"query": question}))
|
|
reply=str(qachain.run(question))
|
|
return reply
|
|
|
|
|
|
|
|
demo = gr.ChatInterface(AI_response, title="Put your files in folder" + docsUrl)
|
|
|
|
if __name__ == "__main__":
|
|
demo.launch(server_name="0.0.0.0", server_port=7860)
|
|
|