81 lines
3.0 KiB
Python
81 lines
3.0 KiB
Python
#! /usr/bin/python3.10
|
|
|
|
import gradio as gr
|
|
import sys
|
|
import os
|
|
import subprocess
|
|
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
|
|
|
|
docsUrl = "/home/user/dev/docs"
|
|
ollamaModel="llama2"
|
|
|
|
def get_ollama_names():
|
|
output = subprocess.check_output(["ollama", "list"])
|
|
lines = output.decode("utf-8").splitlines()
|
|
names = {}
|
|
for line in lines[1:]:
|
|
name = line.split()[0].split(':')[0]
|
|
names[name] = name
|
|
return names
|
|
|
|
names = get_ollama_names()
|
|
|
|
def greet(name):
|
|
global ollamaModel
|
|
ollamaModel=name
|
|
return f"{name}"
|
|
|
|
dropdown = gr.Dropdown(label="Models available", choices=names, value="llama2")
|
|
textbox = gr.Textbox(label="You chose")
|
|
|
|
def AI_response(question, history):
|
|
ollama = Ollama(base_url='http://localhost:11434', model=ollamaModel)
|
|
print(ollamaModel)
|
|
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)
|
|
vectorstore = Chroma.from_documents(documents=all_splits, embedding=GPT4AllEmbeddings(embeddings_chunk_size=1000))
|
|
docs = vectorstore.similarity_search(question)
|
|
len(docs)
|
|
qachain=RetrievalQA.from_chain_type(ollama, retriever=vectorstore.as_retriever())
|
|
reply=str(qachain.run(question))
|
|
return reply
|
|
|
|
|
|
|
|
with gr.Blocks() as demo:
|
|
interface = gr.Interface(fn=greet, inputs=[dropdown], outputs=[textbox], title="Choose a LLM model")
|
|
chat = gr.ChatInterface(AI_response, title="Put your files in folder " + docsUrl)
|
|
|
|
demo.launch(server_name="0.0.0.0", server_port=7860)
|
|
|