mainn.py

import os

from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.vectorstores import FAISS

from langchain_huggingface import HuggingFaceEmbeddings
from transformers import pipeline


# Load files
files = [
    "data/bookkeeping.txt",
    "data/invoice.txt",
    "data/profit_loss.txt",
    "data/balance_sheet.txt",
    "data/cash_flow.txt"
]

documents = []

for file in files:
    loader = TextLoader(file)
    documents.extend(loader.load())


# Split documents
splitter = CharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50
)

docs = splitter.split_documents(documents)


# Create embeddings
embeddings = HuggingFaceEmbeddings()


# Create vector database
db = FAISS.from_documents(docs, embeddings)

retriever = db.as_retriever()


# Load local model
generator = pipeline(
    "text-generation",
    model="google/flan-t5-base"
)

# Ask questions
while True:

    question = input("Ask question: ")

    docs = retriever.invoke(question)

    context = "\n".join([d.page_content for d in docs])

    prompt = f"""
Answer the question based on the following context.

Context:
{context}

Question:
{question}
"""

    result = generator(prompt, max_length=200)

    print("\nAnswer:", result[0]["generated_text"], "\n")