feat: add question generation for imagine la and bem dataset (#144)

navapbc · Dec 2, 2024 · 3803803 · 3803803
1 parent 48c6eab
commit 3803803
Show file tree

Hide file tree

Showing 2 changed files with 141 additions and 12 deletions.
diff --git a/app/notebooks/question_answer_generator/generate_q_a_pairs.py b/app/notebooks/question_answer_generator/generate_q_a_pairs.py
@@ -37,6 +37,7 @@ class QuestionAnswerAttributes(QuestionAnswerPair):
     document_source: str
     document_id: UUID
     chunk_id: Optional[UUID]
+    dataset: str
 
 
 class QuestionAnswerList(BaseModel):
@@ -69,7 +70,10 @@ def generate_question_answer_pairs(llm: str, message: str) -> QuestionAnswerList
 
 
 def process_document_or_chunk(
-    document: Document | Chunk, num_of_chunks: int, llm: str
+    document: Document | Chunk,
+    num_of_chunks: int,
+    llm: str,
+    dataset: str,
 ) -> list[QuestionAnswerAttributes]:
     generated_question_answers = generate_question_answer_pairs(
         llm=llm,
@@ -88,6 +92,7 @@ def process_document_or_chunk(
             question=generated_question_answer.question,
             answer=generated_question_answer.answer,
             chunk_id=None if is_document else document.id,
+            dataset=dataset,
         )
         question_answer_list.append(question_answer_item)