mca183's picture
adding text embedding and vector store
c0e05c3
raw
history blame
1.55 kB
from langchain.document_loaders import HuggingFaceDatasetLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
import gradio as gr
# Load the data
loader = HuggingFaceDatasetLoader("databricks/databricks-dolly-15k", "context")
data = loader.load()
# Document Transformers
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
docs = text_splitter.split_documents(data)
# Text Embedding
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-l6-v2",
model_kwargs={'device':'cpu'},
encode_kwargs={'normalize_embeddings': False}
)
# Set up Vector Stores
db = FAISS.from_documents(docs, embeddings)
# Set up retrievers
retriever = db.as_retriever()
def generate(input):
docs = retriever.get_relevant_documents(input)
return docs[0].page_content
def respond(message, chat_history):
bot_message = generate(message)
chat_history.append((message, bot_message))
return "", chat_history
# Set up the chat interface
with gr.Blocks() as demo:
chatbot = gr.Chatbot(height=240) #just to fit the notebook
msg = gr.Textbox(label="Ask away")
btn = gr.Button("Submit")
clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")
btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) #Press enter to submit
demo.queue().launch()