33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from langchain.embeddings import OpenAIEmbeddings
|
|
from langchain.vectorstores import FAISS
|
|
from langchain.document_loaders import TextLoader
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
from langchain.llms import OpenAI
|
|
from langchain.chains import RetrievalQA
|
|
import os
|
|
|
|
def build_vector_store(directory: str):
|
|
"""Load text files from directory, split into chunks, and build FAISS store."""
|
|
loader = TextLoader(directory, glob='**/*.txt')
|
|
documents = loader.load()
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
docs = splitter.split_documents(documents)
|
|
embeddings = OpenAIEmbeddings()
|
|
vector_store = FAISS.from_documents(docs, embeddings)
|
|
return vector_store
|
|
|
|
def main():
|
|
# Load documents and build vector store
|
|
vector_store = build_vector_store('data')
|
|
# Create RetrievalQA chain
|
|
llm = OpenAI(temperature=0)
|
|
qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
|
|
# Sample query
|
|
query = "What is the capital of France?"
|
|
result = qa_chain.run(query)
|
|
print("Query:", query)
|
|
print("Answer:", result)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|