from langchain.embeddings.openai import OpenAIEmbeddings

from langchain.vectorstores import Chroma

from langchain.llms import OpenAI

from langchain.chains import RetrievalQA

# Đặt khóa API của OpenAI

import os

os.environ[“OPENAI_API_KEY”] = “YOUR_OPENAI_API_KEY”

# Khởi tạo mô hình embedding và LLM

embeddings = OpenAIEmbeddings()

llm = OpenAI(temperature=0)

# Tải tài liệu và chia thành các mảnh

with open(“document.txt”, “r”) as f:

    document = f.read()

chunks = document.split(“\n\n”)  # Chia thành các mảnh dựa trên dòng trống

# Tạo thông tin ngữ cảnh cho mỗi mảnh

contextualized_chunks = []

for chunk in chunks:

    context = llm(f”””

    <document>

    {{Nội dung của toàn bộ tài liệu}}

    </document>

    Đây là mảnh cần được đặt trong toàn bộ tài liệu.

    <chunk>

    {{Nội dung của mảnh}}

    </chunk>

    Hãy mô tả một ngữ cảnh ngắn gọn để định vị mảnh này trong toàn bộ tài liệu nhằm cải thiện khả năng tìm kiếm của nó. Chỉ cung cấp ngữ cảnh ngắn gọn.

    “””)

    contextualized_chunks.append(f”{context} {chunk}”)

# Lưu các mảnh đã có ngữ cảnh vào cơ sở dữ liệu vector

db = Chroma.from_texts(contextualized_chunks, embeddings)

# Tạo chuỗi RetrievalQA

retriever = db.as_retriever()

qa = RetrievalQA.from_chain_type(llm=llm, chain_type=”stuff”, retriever=retriever)

# Nhập truy vấn và lấy câu trả lời

query = “Kết quả quý 2 của công ty như thế nào?”

answer = qa.run(query)

# Xuất câu trả lời

print(answer)