mirror of
https://github.com/tiennm99/test-haystack.git
synced 2026-06-05 20:14:00 +00:00
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
from haystack_integrations.components.generators.ollama import OllamaGenerator
|
|
|
|
from haystack import Pipeline, Document
|
|
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
|
|
from haystack.components.builders.prompt_builder import PromptBuilder
|
|
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
|
import os
|
|
import glob
|
|
|
|
document_store = InMemoryDocumentStore()
|
|
|
|
documents = []
|
|
for txt_file in glob.glob("data/*.txt"):
|
|
with open(txt_file, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
documents.append(Document(content=content))
|
|
|
|
document_store.write_documents(documents)
|
|
|
|
prompt_template = """
|
|
Dựa trên những tài liệu sau, trả lời câu hỏi.
|
|
Tài liệu:
|
|
{% for doc in documents %}
|
|
{{ doc.content }}
|
|
{% endfor %}
|
|
Câu hỏi: {{question}}
|
|
Trả lời:
|
|
"""
|
|
|
|
retriever = InMemoryBM25Retriever(document_store=document_store)
|
|
prompt_builder = PromptBuilder(template=prompt_template)
|
|
llm = OllamaGenerator(url = "http://localhost:11434", model="qwen2.5")
|
|
|
|
rag_pipeline = Pipeline()
|
|
rag_pipeline.add_component("retriever", retriever)
|
|
rag_pipeline.add_component("prompt_builder", prompt_builder)
|
|
rag_pipeline.add_component("llm", llm)
|
|
rag_pipeline.connect("retriever", "prompt_builder.documents")
|
|
rag_pipeline.connect("prompt_builder", "llm")
|
|
|
|
question = "Khan là ai?"
|
|
results = rag_pipeline.run(
|
|
{
|
|
"retriever": {"query": question},
|
|
"prompt_builder": {"question": question},
|
|
}
|
|
)
|
|
|
|
print(results["llm"]["replies"])
|