Optimize a "map reduce_documents chain" in LangChain by storing intermediate results

Question

I have a long document and want to apply different map reduce document chains from LangChain to it. The 'map template' is always identical and can be generated in advance and cached. The various 'reduce prompts' can then be applied to the result of the 'map template' prompt, which is generated only once.

My code is:

from langchain.chat_models import ChatOpenAI
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import ReduceDocumentsChain, MapReduceDocumentsChain
from langchain.document_loaders import TextLoader


nTokens = 4000

docs = TextLoader( "largeTextFile.txt" ).load()

# map prompt which remain unchanged
map_prompt = PromptTemplate.from_template("""
Summarize the following text in a clear and concise way:
TEXT:"{text}"
Brief Summary:
""")
          
llm = ChatOpenAI(temperature=0, model_name = "gpt-3.5-turbo-16k")
map_chain = LLMChain(llm=llm, prompt=map_prompt)

reduce_prompts = ["""
Generate a summary of the following text that includes the following elements:

Text:"{text}"
""", """
You are Markdown writer. Generate a summary of the following text in German:

Text:"{text}"
"""]

# loop over reduce prompts
for promptText in reduce_prompts:

    reduce_chain = LLMChain(llm=llm, prompt=PromptTemplate.from_template( promptText ) )

    combine_documents_chain = StuffDocumentsChain(
        llm_chain=reduce_chain, document_variable_name="text"
    )

    # Combines and iteravely reduces the mapped documents
    reduce_documents_chain = ReduceDocumentsChain(
        combine_documents_chain=combine_documents_chain,
        collapse_documents_chain=combine_documents_chain,
        token_max= nTokens ,
    )

    # map_chain could provide a pre computed output, since docs and map_prompt remain unchanged
    map_reduce_chain = MapReduceDocumentsChain(
        llm_chain=map_chain,
        reduce_documents_chain=reduce_documents_chain,
        document_variable_name="text",
        return_intermediate_steps=False,
    )

    text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
        chunk_size = 1000, chunk_overlap = 0
    )

    split_docs = text_splitter.split_documents(docs)
    text = map_reduce_chain.run(split_docs) 

Any ideas how to implemnent this?

I want to cache the results of the "map reduce" step inside the for loop instead of recalculating them in every run.

Regards, Frank

I tried to implement a custom chain to cache the results and reload if if already computed