-
Notifications
You must be signed in to change notification settings - Fork 0
/
sample_qa.py
45 lines (38 loc) · 1.49 KB
/
sample_qa.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
from langchain.chains import RetrievalQA
from langchain.chains.base import Chain
from langchain.document_loaders import TextLoader
from langchain.indexes.vectorstore import VectorstoreIndexCreator
from langchain.schema.embeddings import Embeddings
from langchain.schema.language_model import BaseLanguageModel
from langchain.text_splitter import TextSplitter
class SampleQA:
"""Define the flow of the model to be adjusted."""
def __init__(
self,
data_path: str,
embedding: Embeddings,
text_splitter: TextSplitter,
llm: BaseLanguageModel,
) -> None:
"""Input the elements necessary for LLM flow The arguments here will be used as a
hyperparameters and optimized.
the arguments are defined by `configs/model/sample.yaml`
"""
self.embedding = embedding
self.text_splitter = text_splitter
self.text_loader = TextLoader(data_path)
self.llm = llm
self.index = VectorstoreIndexCreator(
embedding=self.embedding, text_splitter=self.text_splitter
).from_loaders([self.text_loader])
self.chain = RetrievalQA.from_chain_type(
self.llm,
retriever=self.index.vectorstore.as_retriever(),
return_source_documents=True,
)
def __call__(self, question: str) -> str:
"""Answer the question."""
return self.chain(question)
def get_chain(self) -> Chain:
"""Get langchain chain."""
return self.chain