importtimeimportuuidimportnumpyasnpimporttorchfromlangchain_core.documentsimportDocumentfromlangchain_core.embeddingsimportEmbeddingsfromlangchain_core.language_modelsimportBaseLanguageModelfromlangchain_core.vectorstoresimportVectorStore# --- Assume necessary imports from the package ---fromlangchain_nmretimport ( NeuralMemoryRetriever, TitansNeuralMemoryWrapper, VectorStoreContextualMemory, LightThinkerCompressor, ) # --- Mock/Dummy Components (Replace with actual implementations) ---fromlangchain_community.embeddingsimportHuggingFaceEmbeddingsfromlangchain_community.vectorstoresimportChromafromlangchain_core.language_models.baseimportBaseLanguageModel# For Dummy LLM type hintfromlangchain_core.outputsimportGeneration, LLMResultfromlangchain_core.callbacksimportCallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun, Callbacksfromlangchain_core.prompt_valuesimportPromptValue, StringPromptValuefromlangchain_core.runnablesimportRunnableConfig# Dummy LLM (Replace with OpenAI, HuggingFace, etc.)classDummyRunnable(BaseLanguageModel): def_generate(self, prompts: list[str], stop: list[str] |None=None, run_managers: list[CallbackManagerForLLMRun] |None=None, **kwargs) ->LLMResult: generations= [] fori, promptinenumerate(prompts): text=f"Dummy response to: {prompt.split()[-1]}..."gen= [Generation(text=text)] generations.append(gen) ifrun_managersandrun_managers[i]: run_managers[i].on_llm_end(LLMResult(generations=[gen])) returnLLMResult(generations=generations) asyncdef_agenerate(self, prompts: list[str], stop: list[str] |None=None, run_managers: list[AsyncCallbackManagerForLLMRun] |None=None, **kwargs) ->LLMResult: # Simplified async version for examplereturnself._generate(prompts, stop, None, **kwargs) # Non-async managers for simplicity heredefgenerate_prompt(self, prompts: list[PromptValue], stop: list[str] |None=None, callbacks: Callbacks=None, **kwargs) ->LLMResult: prompt_strings= [str(p) forpinprompts] # Simplified manager handling for dummyreturnself._generate(prompt_strings, stop=stop, **kwargs) asyncdefagenerate_prompt(self, prompts: list[PromptValue], stop: list[str] |None=None, callbacks: Callbacks=None, **kwargs) ->LLMResult: prompt_strings= [str(p) forpinprompts] # Simplified manager handling for dummyreturnawaitself._agenerate(prompt_strings, stop=stop, **kwargs) @propertydef_llm_type(self) ->str: return"dummy"# --- Configuration ---DEVICE="cuda"iftorch.cuda.is_available() else"cpu"EMBEDDING_DIM=384# Example dimension for MiniLM# 1. Embedding Modelembedding_model=HuggingFaceEmbeddings( model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": DEVICE} ) # 2. Vector Storevectorstore=Chroma( collection_name="nmret_readme_example", embedding_function=embedding_model, persist_directory="./chroma_db_readme_example", ) # Add some initial data if neededvectorstore.add_texts( ["Initial context document 1.", "Another piece of information."], metadatas=[{"source": "readme", "memory_id": f"readme_{i}"} foriinrange(2)], ) # 3. Titans Neural Memory Wrappertitans_wrapper=TitansNeuralMemoryWrapper( embedding_dim=EMBEDDING_DIM, device=DEVICE, momentum=True, # Control momentum usage# Pass other NeuralMemory args directly, e.g.:layers=1, heads=2, # Ensure embedding_dim (384) is divisible by heads (2)chunk_size=128, use_accelerated_scan=True# Set to False if assoc-scan not installed ) # 4. Contextual Memory Wrappercontextual_memory=VectorStoreContextualMemory( vectorstore=vectorstore, embedding_model=embedding_model ) # 5. LLMllm=DummyRunnable() # Replace with your actual LLM instance# 6. LightThinker Compressorcompressor=LightThinkerCompressor(llm=llm, embedding_model=embedding_model) # 7. Create the Retrieverretriever=NeuralMemoryRetriever( vectorstore=vectorstore, neural_memory=titans_wrapper, contextual_memory=contextual_memory, compressor=compressor, llm=llm, embedding_model=embedding_model, device=DEVICE, # Configurationreasoning_steps=2, # Number of reasoning loopstop_k_initial=3, # K for initial vector store searchtop_k_contextual=2, # K for contextual memory search per stepcompress_intermediate=True,# Enable LLM thought compressionupdate_memory_on_final=False, # Update Titans memory after each stepupdate_titans_with="docs_and_llm", # Data source for Titans update  ) # --- Run a Query ---query="What is the main topic discussed?"print(f"Invoking retriever with query: '{query}'") results=retriever.invoke(query) print("--- Retrieval Complete ---") print(f"Final Documents Returned: {len(results)}") fori, docinenumerate(results): doc_id=doc.metadata.get("memory_id", "N/A") content_preview=doc.page_content[:100] print(f" - Doc {i}: ID: {doc_id}, Content: {content_preview}...") # Example of saving/loading state (if needed)# state = titans_wrapper.get_state()# # ... save state ...# # ... load state ...# titans_wrapper.load_state(loaded_state)