We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py`
Content:
```
1 """
2 Chain for question-answering against a vector database.
3 
4 Modified from Original Source
5 
6 This code is based on LangChain Ai's langchain, which can be found at
7 https://github.com/langchain-ai/langchain
8 The original code is licensed under the MIT license.
9 """
10 from __future__ import annotations
11 
12 import copy
13 import inspect
14 from typing import Any, Dict, List, Optional
15 
16 from colossalqa.chain.retrieval_qa.load_chain import load_qa_chain
17 from colossalqa.chain.retrieval_qa.stuff import CustomStuffDocumentsChain
18 from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun, Callbacks
19 from langchain.chains.llm import LLMChain
20 from langchain.chains.question_answering.stuff_prompt import PROMPT_SELECTOR
21 from langchain.chains.retrieval_qa.base import BaseRetrievalQA
22 from langchain.prompts import PromptTemplate
23 from langchain.pydantic_v1 import Field
24 from langchain.schema import BaseRetriever, Document
25 from langchain.schema.language_model import BaseLanguageModel
26 
27 class CustomBaseRetrievalQA(BaseRetrievalQA):
28     """Base class for question-answering chains."""
29 
30     @classmethod
31     def from_llm(
32         cls,
33         llm: BaseLanguageModel,
34         prompt: Optional[PromptTemplate] = None,
35         callbacks: Callbacks = None,
36         **kwargs: Any,
37     ) -> BaseRetrievalQA:
38         """Initialize from LLM."""
39         llm_kwargs = kwargs.pop("llm_kwargs", {})
40         _prompt = prompt or PROMPT_SELECTOR.get_prompt(llm)
41         llm_chain = LLMChain(llm=llm, prompt=_prompt, callbacks=callbacks, llm_kwargs=llm_kwargs)
42         document_prompt = kwargs.get(
43             "document_prompt", PromptTemplate(input_variables=["page_content"], template="Context:\n{page_content}")
44         )
45         combine_documents_chain = CustomStuffDocumentsChain(
46             llm_chain=llm_chain,
47             document_variable_name="context",
48             document_prompt=document_prompt,
49             callbacks=callbacks,
50         )
51 
52         return cls(
53             combine_documents_chain=combine_documents_chain,
54             callbacks=callbacks,
55             **kwargs,
56         )
57 
58     @classmethod
59     def from_chain_type(
60         cls,
61         llm: BaseLanguageModel,
62         chain_type: str = "stuff",
63         chain_type_kwargs: Optional[dict] = None,
64         **kwargs: Any,
65     ) -> BaseRetrievalQA:
66         """Load chain from chain type."""
67         llm_kwargs = kwargs.pop("llm_kwargs", {})
68         _chain_type_kwargs = chain_type_kwargs or {}
69         combine_documents_chain = load_qa_chain(llm, chain_type=chain_type, **_chain_type_kwargs, llm_kwargs=llm_kwargs)
70         return cls(combine_documents_chain=combine_documents_chain, **kwargs)
71 
72     def _call(
73         self,
74         inputs: Dict[str, Any],
75         run_manager: Optional[CallbackManagerForChainRun] = None,
76     ) -> Dict[str, Any]:
77         """Run get_relevant_text and llm on input query.
78 
79         If chain has 'return_source_documents' as 'True', returns
80         the retrieved documents as well under the key 'source_documents'.
81 
82         Example:
83         .. code-block:: python
84 
85         res = indexqa({'query': 'This is my query'})
86         answer, docs = res['result'], res['source_documents']
87         """
88         _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
89         question = inputs[self.input_key]
90         accepts_run_manager = "run_manager" in inspect.signature(self._get_docs).parameters
91         if accepts_run_manager:
92             docs = self._get_docs(question, run_manager=_run_manager)
93         else:
94             docs = self._get_docs(question)  # type: ignore[call-arg]
95 
96         kwargs = {
97             k: v
98             for k, v in inputs.items()
99             if k in ["stop", "temperature", "top_k", "top_p", "max_new_tokens", "doc_prefix"]
100         }
101         answers = []
102         if self.combine_documents_chain.memory is not None:
103             buffered_history_backup, summarized_history_temp_backup = copy.deepcopy(
104                 self.combine_documents_chain.memory.buffered_history
105             ), copy.deepcopy(self.combine_documents_chain.memory.summarized_history_temp)
106         else:
107             buffered_history_backup = None
108             summarized_history_temp_backup = None
109 
110         answer = self.combine_documents_chain.run(
111             input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs
112         )
113         if summarized_history_temp_backup is not None and buffered_history_backup is not None:
114             (
115                 self.combine_documents_chain.memory.buffered_history,
116                 self.combine_documents_chain.memory.summarized_history_temp,
117             ) = copy.deepcopy(buffered_history_backup), copy.deepcopy(summarized_history_temp_backup)
118 
119         # if rejection_trigger_keywords is not given, return the response from LLM directly
120         rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])
121         answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) else None
122         if answer is None: 
123             answer = inputs.get('rejection_answer', "抱歉，根据提供的信息无法回答该问题。")
124         if self.combine_documents_chain.memory is not None:
125             self.combine_documents_chain.memory.save_context({"question": question}, {"output": answer})
126 
127         if self.return_source_documents:
128             return {self.output_key: answer, "source_documents": docs}
129         else:
130             return {self.output_key: answer}
131 
132     async def _acall(
133         self,
134         inputs: Dict[str, Any],
135         run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
136     ) -> Dict[str, Any]:
137         """Run get_relevant_text and llm on input query.
138 
139         If chain has 'return_source_documents' as 'True', returns
140         the retrieved documents as well under the key 'source_documents'.
141 
142         Example:
143         .. code-block:: python
144 
145         res = indexqa({'query': 'This is my query'})
146         answer, docs = res['result'], res['source_documents']
147         """
148         _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
149         question = inputs[self.input_key]
150         accepts_run_manager = "run_manager" in inspect.signature(self._aget_docs).parameters
151         if accepts_run_manager:
152             docs = await self._aget_docs(question, run_manager=_run_manager)
153         else:
154             docs = await self._aget_docs(question)  # type: ignore[call-arg]
155         kwargs = {
156             k: v
157             for k, v in inputs.items()
158             if k in ["stop", "temperature", "top_k", "top_p", "max_new_tokens", "doc_prefix"]
159         }
160         answer = await self.combine_documents_chain.arun(
161             input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs
162         )
163         # if rejection_trigger_keywords is not given, return the response from LLM directly
164         rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])
165         answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods)==0 else None
166         if answer is None:
167             answer = inputs.get('rejection_answer', "抱歉，根据提供的信息无法回答该问题。")
168         self.combine_documents_chain.memory.save_context({"question": question}, {"output": answer})
169 
170         if self.return_source_documents:
171             return {self.output_key: answer, "source_documents": docs}
172         else:
173             return {self.output_key: answer}
174 
175 
176 class RetrievalQA(CustomBaseRetrievalQA):
177     """Chain for question-answering against an index.
178 
179     Example:
180         .. code-block:: python
181 
182             from langchain.llms import OpenAI
183             from langchain.chains import RetrievalQA
184             from langchain.faiss import FAISS
185             from langchain.vectorstores.base import VectorStoreRetriever
186             retriever = VectorStoreRetriever(vectorstore=FAISS(...))
187             retrievalQA = RetrievalQA.from_llm(llm=OpenAI(), retriever=retriever)
188 
189     """
190 
191     retriever: BaseRetriever = Field(exclude=True)
192 
193     def _get_docs(
194         self,
195         question: str,
196         *,
197         run_manager: CallbackManagerForChainRun,
198     ) -> List[Document]:
199         """Get docs."""
200         return self.retriever.get_relevant_documents(question, callbacks=run_manager.get_child())
201 
202     async def _aget_docs(
203         self,
204         question: str,
205         *,
206         run_manager: AsyncCallbackManagerForChainRun,
207     ) -> List[Document]:
208         """Get docs."""
209         return await self.retriever.aget_relevant_documents(question, callbacks=run_manager.get_child())
210 
211     @property
212     def _chain_type(self) -> str:
213         """Return the chain type."""
214         return "retrieval_qa"
215
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py b/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py
--- a/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py
+++ b/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py
@@ -24,6 +24,7 @@
 from langchain.schema import BaseRetriever, Document
 from langchain.schema.language_model import BaseLanguageModel
 
+
 class CustomBaseRetrievalQA(BaseRetrievalQA):
     """Base class for question-answering chains."""
 
@@ -98,7 +99,6 @@
             for k, v in inputs.items()
             if k in ["stop", "temperature", "top_k", "top_p", "max_new_tokens", "doc_prefix"]
         }
-        answers = []
         if self.combine_documents_chain.memory is not None:
             buffered_history_backup, summarized_history_temp_backup = copy.deepcopy(
                 self.combine_documents_chain.memory.buffered_history
@@ -117,10 +117,10 @@
             ) = copy.deepcopy(buffered_history_backup), copy.deepcopy(summarized_history_temp_backup)
 
         # if rejection_trigger_keywords is not given, return the response from LLM directly
-        rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])
+        rejection_trigger_keywrods = inputs.get("rejection_trigger_keywrods", [])
         answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) else None
-        if answer is None: 
-            answer = inputs.get('rejection_answer', "抱歉，根据提供的信息无法回答该问题。")
+        if answer is None:
+            answer = inputs.get("rejection_answer", "抱歉，根据提供的信息无法回答该问题。")
         if self.combine_documents_chain.memory is not None:
             self.combine_documents_chain.memory.save_context({"question": question}, {"output": answer})
 
@@ -161,10 +161,14 @@
             input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs
         )
         # if rejection_trigger_keywords is not given, return the response from LLM directly
-        rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])
-        answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods)==0 else None
+        rejection_trigger_keywrods = inputs.get("rejection_trigger_keywrods", [])
+        answer = (
+            answer
+            if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods) == 0
+            else None
+        )
         if answer is None:
-            answer = inputs.get('rejection_answer', "抱歉，根据提供的信息无法回答该问题。")
+            answer = inputs.get("rejection_answer", "抱歉，根据提供的信息无法回答该问题。")
         self.combine_documents_chain.memory.save_context({"question": question}, {"output": answer})
 
         if self.return_source_documents:

{"golden_diff": "diff --git a/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py b/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py\n--- a/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py\n+++ b/applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py\n@@ -24,6 +24,7 @@\n from langchain.schema import BaseRetriever, Document\n from langchain.schema.language_model import BaseLanguageModel\n \n+\n class CustomBaseRetrievalQA(BaseRetrievalQA):\n     \"\"\"Base class for question-answering chains.\"\"\"\n \n@@ -98,7 +99,6 @@\n             for k, v in inputs.items()\n             if k in [\"stop\", \"temperature\", \"top_k\", \"top_p\", \"max_new_tokens\", \"doc_prefix\"]\n         }\n-        answers = []\n         if self.combine_documents_chain.memory is not None:\n             buffered_history_backup, summarized_history_temp_backup = copy.deepcopy(\n                 self.combine_documents_chain.memory.buffered_history\n@@ -117,10 +117,10 @@\n             ) = copy.deepcopy(buffered_history_backup), copy.deepcopy(summarized_history_temp_backup)\n \n         # if rejection_trigger_keywords is not given, return the response from LLM directly\n-        rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])\n+        rejection_trigger_keywrods = inputs.get(\"rejection_trigger_keywrods\", [])\n         answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) else None\n-        if answer is None: \n-            answer = inputs.get('rejection_answer', \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n+        if answer is None:\n+            answer = inputs.get(\"rejection_answer\", \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n         if self.combine_documents_chain.memory is not None:\n             self.combine_documents_chain.memory.save_context({\"question\": question}, {\"output\": answer})\n \n@@ -161,10 +161,14 @@\n             input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs\n         )\n         # if rejection_trigger_keywords is not given, return the response from LLM directly\n-        rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])\n-        answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods)==0 else None\n+        rejection_trigger_keywrods = inputs.get(\"rejection_trigger_keywrods\", [])\n+        answer = (\n+            answer\n+            if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods) == 0\n+            else None\n+        )\n         if answer is None:\n-            answer = inputs.get('rejection_answer', \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n+            answer = inputs.get(\"rejection_answer\", \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n         self.combine_documents_chain.memory.save_context({\"question\": question}, {\"output\": answer})\n \n         if self.return_source_documents:\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "\"\"\"\nChain for question-answering against a vector database.\n\nModified from Original Source\n\nThis code is based on LangChain Ai's langchain, which can be found at\nhttps://github.com/langchain-ai/langchain\nThe original code is licensed under the MIT license.\n\"\"\"\nfrom __future__ import annotations\n\nimport copy\nimport inspect\nfrom typing import Any, Dict, List, Optional\n\nfrom colossalqa.chain.retrieval_qa.load_chain import load_qa_chain\nfrom colossalqa.chain.retrieval_qa.stuff import CustomStuffDocumentsChain\nfrom langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun, Callbacks\nfrom langchain.chains.llm import LLMChain\nfrom langchain.chains.question_answering.stuff_prompt import PROMPT_SELECTOR\nfrom langchain.chains.retrieval_qa.base import BaseRetrievalQA\nfrom langchain.prompts import PromptTemplate\nfrom langchain.pydantic_v1 import Field\nfrom langchain.schema import BaseRetriever, Document\nfrom langchain.schema.language_model import BaseLanguageModel\n\nclass CustomBaseRetrievalQA(BaseRetrievalQA):\n    \"\"\"Base class for question-answering chains.\"\"\"\n\n    @classmethod\n    def from_llm(\n        cls,\n        llm: BaseLanguageModel,\n        prompt: Optional[PromptTemplate] = None,\n        callbacks: Callbacks = None,\n        **kwargs: Any,\n    ) -> BaseRetrievalQA:\n        \"\"\"Initialize from LLM.\"\"\"\n        llm_kwargs = kwargs.pop(\"llm_kwargs\", {})\n        _prompt = prompt or PROMPT_SELECTOR.get_prompt(llm)\n        llm_chain = LLMChain(llm=llm, prompt=_prompt, callbacks=callbacks, llm_kwargs=llm_kwargs)\n        document_prompt = kwargs.get(\n            \"document_prompt\", PromptTemplate(input_variables=[\"page_content\"], template=\"Context:\\n{page_content}\")\n        )\n        combine_documents_chain = CustomStuffDocumentsChain(\n            llm_chain=llm_chain,\n            document_variable_name=\"context\",\n            document_prompt=document_prompt,\n            callbacks=callbacks,\n        )\n\n        return cls(\n            combine_documents_chain=combine_documents_chain,\n            callbacks=callbacks,\n            **kwargs,\n        )\n\n    @classmethod\n    def from_chain_type(\n        cls,\n        llm: BaseLanguageModel,\n        chain_type: str = \"stuff\",\n        chain_type_kwargs: Optional[dict] = None,\n        **kwargs: Any,\n    ) -> BaseRetrievalQA:\n        \"\"\"Load chain from chain type.\"\"\"\n        llm_kwargs = kwargs.pop(\"llm_kwargs\", {})\n        _chain_type_kwargs = chain_type_kwargs or {}\n        combine_documents_chain = load_qa_chain(llm, chain_type=chain_type, **_chain_type_kwargs, llm_kwargs=llm_kwargs)\n        return cls(combine_documents_chain=combine_documents_chain, **kwargs)\n\n    def _call(\n        self,\n        inputs: Dict[str, Any],\n        run_manager: Optional[CallbackManagerForChainRun] = None,\n    ) -> Dict[str, Any]:\n        \"\"\"Run get_relevant_text and llm on input query.\n\n        If chain has 'return_source_documents' as 'True', returns\n        the retrieved documents as well under the key 'source_documents'.\n\n        Example:\n        .. code-block:: python\n\n        res = indexqa({'query': 'This is my query'})\n        answer, docs = res['result'], res['source_documents']\n        \"\"\"\n        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()\n        question = inputs[self.input_key]\n        accepts_run_manager = \"run_manager\" in inspect.signature(self._get_docs).parameters\n        if accepts_run_manager:\n            docs = self._get_docs(question, run_manager=_run_manager)\n        else:\n            docs = self._get_docs(question)  # type: ignore[call-arg]\n\n        kwargs = {\n            k: v\n            for k, v in inputs.items()\n            if k in [\"stop\", \"temperature\", \"top_k\", \"top_p\", \"max_new_tokens\", \"doc_prefix\"]\n        }\n        answers = []\n        if self.combine_documents_chain.memory is not None:\n            buffered_history_backup, summarized_history_temp_backup = copy.deepcopy(\n                self.combine_documents_chain.memory.buffered_history\n            ), copy.deepcopy(self.combine_documents_chain.memory.summarized_history_temp)\n        else:\n            buffered_history_backup = None\n            summarized_history_temp_backup = None\n\n        answer = self.combine_documents_chain.run(\n            input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs\n        )\n        if summarized_history_temp_backup is not None and buffered_history_backup is not None:\n            (\n                self.combine_documents_chain.memory.buffered_history,\n                self.combine_documents_chain.memory.summarized_history_temp,\n            ) = copy.deepcopy(buffered_history_backup), copy.deepcopy(summarized_history_temp_backup)\n\n        # if rejection_trigger_keywords is not given, return the response from LLM directly\n        rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])\n        answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) else None\n        if answer is None: \n            answer = inputs.get('rejection_answer', \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n        if self.combine_documents_chain.memory is not None:\n            self.combine_documents_chain.memory.save_context({\"question\": question}, {\"output\": answer})\n\n        if self.return_source_documents:\n            return {self.output_key: answer, \"source_documents\": docs}\n        else:\n            return {self.output_key: answer}\n\n    async def _acall(\n        self,\n        inputs: Dict[str, Any],\n        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,\n    ) -> Dict[str, Any]:\n        \"\"\"Run get_relevant_text and llm on input query.\n\n        If chain has 'return_source_documents' as 'True', returns\n        the retrieved documents as well under the key 'source_documents'.\n\n        Example:\n        .. code-block:: python\n\n        res = indexqa({'query': 'This is my query'})\n        answer, docs = res['result'], res['source_documents']\n        \"\"\"\n        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()\n        question = inputs[self.input_key]\n        accepts_run_manager = \"run_manager\" in inspect.signature(self._aget_docs).parameters\n        if accepts_run_manager:\n            docs = await self._aget_docs(question, run_manager=_run_manager)\n        else:\n            docs = await self._aget_docs(question)  # type: ignore[call-arg]\n        kwargs = {\n            k: v\n            for k, v in inputs.items()\n            if k in [\"stop\", \"temperature\", \"top_k\", \"top_p\", \"max_new_tokens\", \"doc_prefix\"]\n        }\n        answer = await self.combine_documents_chain.arun(\n            input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs\n        )\n        # if rejection_trigger_keywords is not given, return the response from LLM directly\n        rejection_trigger_keywrods = inputs.get('rejection_trigger_keywrods', [])\n        answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods)==0 else None\n        if answer is None:\n            answer = inputs.get('rejection_answer', \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n        self.combine_documents_chain.memory.save_context({\"question\": question}, {\"output\": answer})\n\n        if self.return_source_documents:\n            return {self.output_key: answer, \"source_documents\": docs}\n        else:\n            return {self.output_key: answer}\n\n\nclass RetrievalQA(CustomBaseRetrievalQA):\n    \"\"\"Chain for question-answering against an index.\n\n    Example:\n        .. code-block:: python\n\n            from langchain.llms import OpenAI\n            from langchain.chains import RetrievalQA\n            from langchain.faiss import FAISS\n            from langchain.vectorstores.base import VectorStoreRetriever\n            retriever = VectorStoreRetriever(vectorstore=FAISS(...))\n            retrievalQA = RetrievalQA.from_llm(llm=OpenAI(), retriever=retriever)\n\n    \"\"\"\n\n    retriever: BaseRetriever = Field(exclude=True)\n\n    def _get_docs(\n        self,\n        question: str,\n        *,\n        run_manager: CallbackManagerForChainRun,\n    ) -> List[Document]:\n        \"\"\"Get docs.\"\"\"\n        return self.retriever.get_relevant_documents(question, callbacks=run_manager.get_child())\n\n    async def _aget_docs(\n        self,\n        question: str,\n        *,\n        run_manager: AsyncCallbackManagerForChainRun,\n    ) -> List[Document]:\n        \"\"\"Get docs.\"\"\"\n        return await self.retriever.aget_relevant_documents(question, callbacks=run_manager.get_child())\n\n    @property\n    def _chain_type(self) -> str:\n        \"\"\"Return the chain type.\"\"\"\n        return \"retrieval_qa\"\n", "path": "applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py"}], "after_files": [{"content": "\"\"\"\nChain for question-answering against a vector database.\n\nModified from Original Source\n\nThis code is based on LangChain Ai's langchain, which can be found at\nhttps://github.com/langchain-ai/langchain\nThe original code is licensed under the MIT license.\n\"\"\"\nfrom __future__ import annotations\n\nimport copy\nimport inspect\nfrom typing import Any, Dict, List, Optional\n\nfrom colossalqa.chain.retrieval_qa.load_chain import load_qa_chain\nfrom colossalqa.chain.retrieval_qa.stuff import CustomStuffDocumentsChain\nfrom langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun, Callbacks\nfrom langchain.chains.llm import LLMChain\nfrom langchain.chains.question_answering.stuff_prompt import PROMPT_SELECTOR\nfrom langchain.chains.retrieval_qa.base import BaseRetrievalQA\nfrom langchain.prompts import PromptTemplate\nfrom langchain.pydantic_v1 import Field\nfrom langchain.schema import BaseRetriever, Document\nfrom langchain.schema.language_model import BaseLanguageModel\n\n\nclass CustomBaseRetrievalQA(BaseRetrievalQA):\n    \"\"\"Base class for question-answering chains.\"\"\"\n\n    @classmethod\n    def from_llm(\n        cls,\n        llm: BaseLanguageModel,\n        prompt: Optional[PromptTemplate] = None,\n        callbacks: Callbacks = None,\n        **kwargs: Any,\n    ) -> BaseRetrievalQA:\n        \"\"\"Initialize from LLM.\"\"\"\n        llm_kwargs = kwargs.pop(\"llm_kwargs\", {})\n        _prompt = prompt or PROMPT_SELECTOR.get_prompt(llm)\n        llm_chain = LLMChain(llm=llm, prompt=_prompt, callbacks=callbacks, llm_kwargs=llm_kwargs)\n        document_prompt = kwargs.get(\n            \"document_prompt\", PromptTemplate(input_variables=[\"page_content\"], template=\"Context:\\n{page_content}\")\n        )\n        combine_documents_chain = CustomStuffDocumentsChain(\n            llm_chain=llm_chain,\n            document_variable_name=\"context\",\n            document_prompt=document_prompt,\n            callbacks=callbacks,\n        )\n\n        return cls(\n            combine_documents_chain=combine_documents_chain,\n            callbacks=callbacks,\n            **kwargs,\n        )\n\n    @classmethod\n    def from_chain_type(\n        cls,\n        llm: BaseLanguageModel,\n        chain_type: str = \"stuff\",\n        chain_type_kwargs: Optional[dict] = None,\n        **kwargs: Any,\n    ) -> BaseRetrievalQA:\n        \"\"\"Load chain from chain type.\"\"\"\n        llm_kwargs = kwargs.pop(\"llm_kwargs\", {})\n        _chain_type_kwargs = chain_type_kwargs or {}\n        combine_documents_chain = load_qa_chain(llm, chain_type=chain_type, **_chain_type_kwargs, llm_kwargs=llm_kwargs)\n        return cls(combine_documents_chain=combine_documents_chain, **kwargs)\n\n    def _call(\n        self,\n        inputs: Dict[str, Any],\n        run_manager: Optional[CallbackManagerForChainRun] = None,\n    ) -> Dict[str, Any]:\n        \"\"\"Run get_relevant_text and llm on input query.\n\n        If chain has 'return_source_documents' as 'True', returns\n        the retrieved documents as well under the key 'source_documents'.\n\n        Example:\n        .. code-block:: python\n\n        res = indexqa({'query': 'This is my query'})\n        answer, docs = res['result'], res['source_documents']\n        \"\"\"\n        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()\n        question = inputs[self.input_key]\n        accepts_run_manager = \"run_manager\" in inspect.signature(self._get_docs).parameters\n        if accepts_run_manager:\n            docs = self._get_docs(question, run_manager=_run_manager)\n        else:\n            docs = self._get_docs(question)  # type: ignore[call-arg]\n\n        kwargs = {\n            k: v\n            for k, v in inputs.items()\n            if k in [\"stop\", \"temperature\", \"top_k\", \"top_p\", \"max_new_tokens\", \"doc_prefix\"]\n        }\n        if self.combine_documents_chain.memory is not None:\n            buffered_history_backup, summarized_history_temp_backup = copy.deepcopy(\n                self.combine_documents_chain.memory.buffered_history\n            ), copy.deepcopy(self.combine_documents_chain.memory.summarized_history_temp)\n        else:\n            buffered_history_backup = None\n            summarized_history_temp_backup = None\n\n        answer = self.combine_documents_chain.run(\n            input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs\n        )\n        if summarized_history_temp_backup is not None and buffered_history_backup is not None:\n            (\n                self.combine_documents_chain.memory.buffered_history,\n                self.combine_documents_chain.memory.summarized_history_temp,\n            ) = copy.deepcopy(buffered_history_backup), copy.deepcopy(summarized_history_temp_backup)\n\n        # if rejection_trigger_keywords is not given, return the response from LLM directly\n        rejection_trigger_keywrods = inputs.get(\"rejection_trigger_keywrods\", [])\n        answer = answer if all([rej not in answer for rej in rejection_trigger_keywrods]) else None\n        if answer is None:\n            answer = inputs.get(\"rejection_answer\", \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n        if self.combine_documents_chain.memory is not None:\n            self.combine_documents_chain.memory.save_context({\"question\": question}, {\"output\": answer})\n\n        if self.return_source_documents:\n            return {self.output_key: answer, \"source_documents\": docs}\n        else:\n            return {self.output_key: answer}\n\n    async def _acall(\n        self,\n        inputs: Dict[str, Any],\n        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,\n    ) -> Dict[str, Any]:\n        \"\"\"Run get_relevant_text and llm on input query.\n\n        If chain has 'return_source_documents' as 'True', returns\n        the retrieved documents as well under the key 'source_documents'.\n\n        Example:\n        .. code-block:: python\n\n        res = indexqa({'query': 'This is my query'})\n        answer, docs = res['result'], res['source_documents']\n        \"\"\"\n        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()\n        question = inputs[self.input_key]\n        accepts_run_manager = \"run_manager\" in inspect.signature(self._aget_docs).parameters\n        if accepts_run_manager:\n            docs = await self._aget_docs(question, run_manager=_run_manager)\n        else:\n            docs = await self._aget_docs(question)  # type: ignore[call-arg]\n        kwargs = {\n            k: v\n            for k, v in inputs.items()\n            if k in [\"stop\", \"temperature\", \"top_k\", \"top_p\", \"max_new_tokens\", \"doc_prefix\"]\n        }\n        answer = await self.combine_documents_chain.arun(\n            input_documents=docs, question=question, callbacks=_run_manager.get_child(), **kwargs\n        )\n        # if rejection_trigger_keywords is not given, return the response from LLM directly\n        rejection_trigger_keywrods = inputs.get(\"rejection_trigger_keywrods\", [])\n        answer = (\n            answer\n            if all([rej not in answer for rej in rejection_trigger_keywrods]) or len(rejection_trigger_keywrods) == 0\n            else None\n        )\n        if answer is None:\n            answer = inputs.get(\"rejection_answer\", \"\u62b1\u6b49\uff0c\u6839\u636e\u63d0\u4f9b\u7684\u4fe1\u606f\u65e0\u6cd5\u56de\u7b54\u8be5\u95ee\u9898\u3002\")\n        self.combine_documents_chain.memory.save_context({\"question\": question}, {\"output\": answer})\n\n        if self.return_source_documents:\n            return {self.output_key: answer, \"source_documents\": docs}\n        else:\n            return {self.output_key: answer}\n\n\nclass RetrievalQA(CustomBaseRetrievalQA):\n    \"\"\"Chain for question-answering against an index.\n\n    Example:\n        .. code-block:: python\n\n            from langchain.llms import OpenAI\n            from langchain.chains import RetrievalQA\n            from langchain.faiss import FAISS\n            from langchain.vectorstores.base import VectorStoreRetriever\n            retriever = VectorStoreRetriever(vectorstore=FAISS(...))\n            retrievalQA = RetrievalQA.from_llm(llm=OpenAI(), retriever=retriever)\n\n    \"\"\"\n\n    retriever: BaseRetriever = Field(exclude=True)\n\n    def _get_docs(\n        self,\n        question: str,\n        *,\n        run_manager: CallbackManagerForChainRun,\n    ) -> List[Document]:\n        \"\"\"Get docs.\"\"\"\n        return self.retriever.get_relevant_documents(question, callbacks=run_manager.get_child())\n\n    async def _aget_docs(\n        self,\n        question: str,\n        *,\n        run_manager: AsyncCallbackManagerForChainRun,\n    ) -> List[Document]:\n        \"\"\"Get docs.\"\"\"\n        return await self.retriever.aget_relevant_documents(question, callbacks=run_manager.get_child())\n\n    @property\n    def _chain_type(self) -> str:\n        \"\"\"Return the chain type.\"\"\"\n        return \"retrieval_qa\"\n", "path": "applications/ColossalQA/colossalqa/chain/retrieval_qa/base.py"}]}