feat: extend QA generation types (#126)

Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
DS4SD · Aug 7, 2023 · 8900d17 · 8900d17
1 parent b93ed58
commit 8900d17
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 6 deletions.
diff --git a/deepsearch/model/examples/dummy_qa_generator/model.py b/deepsearch/model/examples/dummy_qa_generator/model.py
@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 from deepsearch.model.base.types import Kind
 from deepsearch.model.kinds.qagen.model import BaseQAGenerator
@@ -18,7 +18,7 @@ def get_qagen_config(self) -> QAGenConfig:
         return self._config
 
     def generate_answers(
-        self, texts: List[Tuple[List[str], str]]
+        self, texts: List[Tuple[List[Dict], str]]
     ) -> GenerateAnswersOutput:
         """Just answers with the question itself.
         Args:

diff --git a/deepsearch/model/kinds/qagen/controller.py b/deepsearch/model/kinds/qagen/controller.py
@@ -39,7 +39,10 @@ def dispatch_predict(self, spec: CtrlPredInput) -> CtrlPredOutput:
         if isinstance(spec, QAGenReqSpec):
             gen_answers = spec.generateAnswers
             answers = self._model.generate_answers(
-                [(c, q) for c, q in zip(gen_answers.contexts, gen_answers.questions)]
+                texts=[
+                    ([ctx_entry.dict() for ctx_entry in ctx_list], q)
+                    for ctx_list, q in zip(gen_answers.contexts, gen_answers.questions)
+                ],
             )
             return QAGenCtrlPredOutput(
                 answers=answers,

diff --git a/deepsearch/model/kinds/qagen/model.py b/deepsearch/model/kinds/qagen/model.py
@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 from deepsearch.model.base.model import BaseDSModel
 from deepsearch.model.base.types import BaseModelConfig
@@ -9,7 +9,7 @@
 class BaseQAGenerator(BaseDSModel):
     @abstractmethod
     def generate_answers(
-        self, texts: List[Tuple[List[str], str]]
+        self, texts: List[Tuple[List[Dict], str]]
     ) -> GenerateAnswersOutput:
         raise NotImplementedError()
 

diff --git a/deepsearch/model/kinds/qagen/types.py b/deepsearch/model/kinds/qagen/types.py
@@ -11,8 +11,14 @@
 )
 
 
+class ContextEntry(StrictModel):
+    text: str
+    type: str
+    representation_type: str
+
+
 class GenerateAnswers(StrictModel):
-    contexts: List[List[str]]
+    contexts: List[List[ContextEntry]]
     questions: List[str]
 
     @root_validator