File size: 4,538 Bytes
1afde00 11a4786 9dee257 1afde00 af5492a df9ff7b 4879c2c df9ff7b 0dc9dbf 11a4786 af5492a df9ff7b 3cd0964 fe8a215 3cd0964 fe8a215 3cd0964 fe8a215 3cd0964 af5492a 423c62e af5492a 423c62e af5492a 423c62e fe8a215 423c62e fe8a215 af5492a 3cd0964 af5492a 3cd0964 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import os
import json
import re
import sys
import gradio as gr
from huggingface_hub import InferenceClient
from langchain_huggingface import HuggingFaceEmbeddings
#from chromadb.utils import embedding_functions
#from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.document_loaders import PyPDFLoader
from fastapi.encoders import jsonable_encoder
"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
# Select which embeddings we want to use
#embeddings = OpenAIEmbeddings()
#embeddings = SentenceTransformerEmbeddings(model_name="nomic-ai/nomic-embed-text-v1", model_kwargs={"trust_remote_code":True})
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
ABS_PATH = os.path.dirname(os.path.abspath(__file__))
DB_DIR = os.path.join(ABS_PATH, "db")
vectorstore = None
def replace_newlines_and_spaces(text):
# Replace all newline characters with spaces
text = text.replace("\n", " ")
# Replace multiple spaces with a single space
text = re.sub(r'\s+', ' ', text)
return text
def get_documents():
return PyPDFLoader("AI-smart-water-management-systems.pdf").load()
def init_chromadb():
# Delete existing index directory and recreate the directory
if os.path.exists(DB_DIR):
import shutil
shutil.rmtree(DB_DIR, ignore_errors=True)
os.mkdir(DB_DIR)
documents = []
for num, doc in enumerate(get_documents()):
doc.page_content = replace_newlines_and_spaces(doc.page_content)
documents.append(doc)
# Split the documents into chunks
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
#query_chromadb()
# Create the vectorestore to use as the index
vectorstore = Chroma.from_documents(texts, embeddings, persist_directory=DB_DIR)
vectorstore.persist()
print("vectorstore::", vectorstore)
def query_chromadb(ASK):
if not os.path.exists(DB_DIR):
raise Exception(f"{DB_DIR} does not exist, nothing can be queried")
# Load Vector store from local disk
vectorstore = Chroma(persist_directory=DB_DIR, embedding_function=embeddings)
result = vectorstore.similarity_search_with_score(query=ASK, k=4)
jsonable_result = jsonable_encoder(result)
print("Json pdf response ::", json.dumps(jsonable_result, indent=2))
#return json.dumps(jsonable_result, indent=2)
def respond(
message,
history: list[tuple[str, str]],
system_message,
max_tokens,
temperature,
top_p,
):
messages = [{"role": "system", "content": system_message}]
for val in history:
if val[0]:
messages.append({"role": "user", "content": val[0]})
if val[1]:
messages.append({"role": "assistant", "content": val[1]})
messages.append({"role": "user", "content": message})
response = ""
for message in client.chat_completion(
messages,
max_tokens=max_tokens,
stream=True,
temperature=temperature,
top_p=top_p,
):
print ("**message :: ",message)
token = message.choices[0].delta.content
print ("**token :: ",token)
response += token
print ("**response :: ",response)
yield response
print ("**query_chromadb::",query_chromadb("how could an AI be used in smart water management systems?"))
#yield query_chromadb(message)
"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
respond,
additional_inputs=[
gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.95,
step=0.05,
label="Top-p (nucleus sampling)",
),
],
)
def main():
init_chromadb()
demo.launch()
if __name__ == "__main__":
main()
#demo.launch()
|