Create initialize.py

initialize.py

@@ -0,0 +1,55 @@
+from haystack import Document
+from haystack.document_stores.in_memory import InMemoryDocumentStore
+from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
+from haystack.components.builders import PromptBuilder
+from haystack.components.generators.hugging_face_local import HuggingFaceLocalGenerator
+from haystack.pipeline import Pipeline
+
+def init_doc_store(path, files):
+    docs = []
+    for file in files:
+        with open(path + '/' + file, 'r') as f:
+            content = f.read()
+            docs.append(Document(content=content, meta={'name':file}))
+
+    document_store = InMemoryDocumentStore()
+    document_store.write_documents(docs)
+    return document_store
+
+def define_components(document_store):
+    retriever = InMemoryBM25Retriever(document_store, top_k=3)
+    
+    template = """
+    Given the following information, answer the question.
+    
+    Context:
+    {% for document in documents %}
+        {{ document.content }}
+    {% endfor %}
+    
+    Question: {{question}}
+    Answer:
+    """
+    prompt_builder = PromptBuilder(template=template)
+    
+    generator = HuggingFaceLocalGenerator(model="gpt2",
+                                          task="text-generation",
+                                        #   device='cuda',
+                                          generation_kwargs={
+                                            "max_new_tokens": 100,
+                                            "temperature": 0.9,
+                                            })
+    generator.warm_up()
+    return retreiver, prompt_builder, generator
+
+def define_pipeline(retreiver, prompt_builder, generator):
+    basic_rag_pipeline = Pipeline()
+    
+    basic_rag_pipeline.add_component("retriever", retriever)
+    basic_rag_pipeline.add_component("prompt_builder", prompt_builder)
+    basic_rag_pipeline.add_component("llm", generator)
+    
+    basic_rag_pipeline.connect("retriever", "prompt_builder.documents")
+    basic_rag_pipeline.connect("prompt_builder", "llm")
+
+    return basic_rag_pipeline