2023-05-12 20:20:07 +00:00
|
|
|
|
|
|
|
import os
|
|
|
|
from langchain.vectorstores import Chroma
|
|
|
|
from langchain.embeddings import OpenAIEmbeddings
|
2023-05-18 20:50:21 +00:00
|
|
|
from langchain.chat_models import ChatOpenAI
|
|
|
|
from langchain.chains import RetrievalQA
|
|
|
|
from langchain.vectorstores.base import VectorStoreRetriever
|
2023-05-12 20:20:07 +00:00
|
|
|
|
|
|
|
base_path = os.environ.get('OPENAI_API_BASE', 'http://localhost:8080/v1')
|
|
|
|
|
|
|
|
# Load and process the text
|
2023-07-27 16:41:53 +00:00
|
|
|
embedding = OpenAIEmbeddings(model="text-embedding-ada-002", openai_api_base=base_path)
|
2023-05-13 09:16:56 +00:00
|
|
|
persist_directory = 'db'
|
2023-05-12 20:20:07 +00:00
|
|
|
|
|
|
|
# Now we can load the persisted database from disk, and use it as normal.
|
2023-05-18 20:50:21 +00:00
|
|
|
llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_base=base_path)
|
2023-05-12 20:20:07 +00:00
|
|
|
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)
|
2023-05-18 20:50:21 +00:00
|
|
|
retriever = VectorStoreRetriever(vectorstore=vectordb)
|
|
|
|
qa = RetrievalQA.from_llm(llm=llm, retriever=retriever)
|
2023-05-12 20:20:07 +00:00
|
|
|
|
|
|
|
query = "What the president said about taxes ?"
|
|
|
|
print(qa.run(query))
|
|
|
|
|