Update initialize.py

initialize.py

@@ -1,4 +1,4 @@
-from haystack import Document
+from haystack import Document, Secret
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
 from haystack.components.builders import PromptBuilder
@@ -16,7 +16,7 @@ def init_doc_store(path, files):
     document_store.write_documents(docs)
     return document_store
 
-def define_components(document_store):
+def define_components(document_store, api_key):
     retriever = InMemoryBM25Retriever(document_store, top_k=3)
     
     template = """
@@ -32,13 +32,10 @@ def define_components(document_store):
     """
     prompt_builder = PromptBuilder(template=template)
     
-    generator = HuggingFaceLocalGenerator(model="gpt2",
-                                          task="text-generation",
-                                        #   device='cuda',
-                                          generation_kwargs={
-                                            "max_new_tokens": 100,
-                                            "temperature": 0.9,
-                                            })
+    generator = HuggingFaceTGIGenerator(
+        model="mistralai/Mistral-7B-v0.1", 
+        token=Secret.from_token(api_key)
+    )
     generator.warm_up()
     return retreiver, prompt_builder, generator