Implementation:Hpcaitech ColossalAI Load QA Chain

Knowledge Sources	Hpcaitech_ColossalAI
Domains	NLP, Question Answering, Chain Loading
Last Updated	2026-02-09 00:00 GMT

Overview

load_qa_chain is a factory function that loads and returns a question answering chain for combining documents, currently supporting the "stuff" chain type via CustomStuffDocumentsChain.

Description

This module provides the load_qa_chain function which acts as a dispatcher for creating document-combining QA chains based on the specified chain type. It also includes the internal _load_stuff_chain helper that constructs a CustomStuffDocumentsChain with an LLMChain and appropriate prompt configuration. The code is modified from LangChain's original implementation and supports passing additional LLM keyword arguments such as generation parameters.

Usage

Use load_qa_chain when building a retrieval-based question answering pipeline in ColossalQA where retrieved documents need to be combined and fed into an LLM for answer generation. This is typically called as part of a retrieval QA workflow.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalQA/colossalqa/chain/retrieval_qa/load_chain.py
• Lines: 1-88

Signature

def load_qa_chain(
    llm: BaseLanguageModel,
    chain_type: str = "stuff",
    verbose: Optional[bool] = None,
    callback_manager: Optional[BaseCallbackManager] = None,
    **kwargs: Any,
) -> BaseCombineDocumentsChain:
    ...

def _load_stuff_chain(
    llm: BaseLanguageModel,
    prompt: Optional[BasePromptTemplate] = None,
    document_variable_name: str = "context",
    verbose: Optional[bool] = None,
    callback_manager: Optional[BaseCallbackManager] = None,
    callbacks: Callbacks = None,
    **kwargs: Any,
) -> CustomStuffDocumentsChain:
    ...

Import

from colossalqa.chain.retrieval_qa.load_chain import load_qa_chain

I/O Contract

Inputs

Name	Type	Required	Description
llm	BaseLanguageModel	Yes	The language model to use in the QA chain
chain_type	str	No	Type of document combining chain to use; must be one of the supported types (default: "stuff")
verbose	Optional[bool]	No	Whether chains should run in verbose mode
callback_manager	Optional[BaseCallbackManager]	No	Callback manager for the chain
prompt	Optional[BasePromptTemplate]	No	Custom prompt template (passed via kwargs to _load_stuff_chain)
document_variable_name	str	No	The variable name for the combined document string in the prompt (default: "context")
llm_kwargs	dict	No	Additional keyword arguments passed to the LLMChain (passed via kwargs)

Outputs

Name	Type	Description
return	BaseCombineDocumentsChain	A configured document-combining chain ready for question answering

Usage Examples

from colossalqa.chain.retrieval_qa.load_chain import load_qa_chain
from langchain.llms import OpenAI

llm = OpenAI()

# Load a stuff-type QA chain
qa_chain = load_qa_chain(
    llm=llm,
    chain_type="stuff",
    verbose=True,
)

# Load with custom LLM generation parameters
qa_chain = load_qa_chain(
    llm=llm,
    chain_type="stuff",
    llm_kwargs={"max_new_tokens": 200, "temperature": 0.7},
)

Related Pages

• Environment:Hpcaitech_ColossalAI_ColossalQA_RAG_Environment