code: """
  This module provides an implementation for generating question data from documents.
  Supported types of document sources include:
      - plain text
      - unstructured files: Text, PDF, PowerPoint, HTML, Images, 
          Excel spreadsheets, Word documents, Markdown, etc.
      - documents from Google Drive (provide file id).
  Currently support only one document a time.
  """
  import ast
  import asyncio
  import csv
  import os
  import pickle
  import re
  from typing import Any, Dict, Iterator, List
  
  from langchain.document_loaders import GoogleDriveLoader, UnstructuredFileLoader
  from langchain.schema import Document
  from langchain.text_splitter import RecursiveCharacterTextSplitter
  from tqdm import tqdm
  
  from yival.common import utils
  from yival.common.model_utils import llm_completion
  from yival.data_generators.base_data_generator import BaseDataGenerator
  from yival.schemas.common_structures import InputData
  from yival.schemas.data_generator_configs import DocumentDataGeneratorConfig
  from yival.schemas.model_configs import Request
  
  PROMPT_TEMPLATE = """
      Context information is below.
      ---------------------
      {CONTEXT}
      ---------------------
      Please do not introduce priori knowledge, 
      only consider the content of the previous context information, 
      generate 5 questions based on the following query. 
      Answer ONLY a python list containing all the questions generated.
      Keep your output crisp, with only a '[]' bracketed list.
      {QUERY}
  """
  
  
  class DocumentDataGenerator(BaseDataGenerator):
      config: DocumentDataGeneratorConfig
      default_config: DocumentDataGeneratorConfig = DocumentDataGeneratorConfig(
          prompt=PROMPT_TEMPLATE,
          document="",
          source="text",
          num_questions_per_chunk=5,
          text_question_template=None,
          document_chunk_size=512,
          number_of_examples=1,
          question_gen_query=f"You are a Teacher/Professor. Your task is to setup \
                          5 questions for an upcoming quiz/examination. The questions \
                          should be diverse in nature across the document. Restrict \
                          the questions to the context information provided."
      )
  
      def __init__(self, config: DocumentDataGeneratorConfig):
          super().__init__(config)
          self.config = config
  
      def load_document(self, source: str, document: str):
          if source == 'text':
              doc = Document(page_content=document)
              return doc
          elif source == 'file':
              file_loader = UnstructuredFileLoader(document)
              docs = file_loader.load()
              doc = docs[0]
              return doc
          elif source == 'drive':
              drive_loader = GoogleDriveLoader(file_ids=[document])
              docs = drive_loader.load()
              doc = docs[0]
              return doc
          else:
              return None
  
      def get_doc_context(self, doc: Document, chunk_size: int) -> List[str]:
          # Split Document into chunks
          splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size)
          splits = splitter.split_documents([doc])
  
          # Generate contexts from splits
          contexts = [str(split.page_content) for split in splits]
          return contexts
  
      def prepare_messages(self) -> List[Dict[str, Any]]:
          """Prepare the messages for GPT API based on configurations."""
          if not self.config.prompt:
              self.config.prompt = PROMPT_TEMPLATE
          document = self.load_document(self.config.source, self.config.document)
          if document:
              contexts = self.get_doc_context(document, self.config.chunk_size)
          else:
              raise TypeError
  
          contents = []
          for context in contexts:
              content = "Context information is below.\n---------------------\n\n" + context + "\n"
              content = content + "---------------------\nPlease do not introduce priori knowledge,\n"
              content = content + "only consider the content of the previous context information,\n generate "
              content = content + str(
                  self.config.num_questions_per_chunk
              ) + " questions based on the following query."
              content = content + "Answer ONLY a python list containing all the questions generated.\n"
              content = content + "Context information is below.\n---------------------\n\n"
              content = content + "Keep your output crisp, with only a '[]' bracketed list.\n"
              content = content + self.config.question_gen_query + "\n"
              if self.config.text_question_template:
                  content = content + "Please generate the questions according to the following template:\n" + self.config.text_question_template + "\n"
              contents.append(content)
          return [{"role": "user", "content": content} for content in contents]
  
      def process_output(
          self, output_content: str, all_data: List[InputData],
          chunk: List[InputData]
      ):
          """Process the output from GPT API and update data lists."""
          output_ls = eval(output_content)
  
          input_data_instance = InputData(
              example_id=super().generate_example_id(output_content),
              content={"data": output_ls}
          )
          all_data.append(input_data_instance)
          chunk.append(input_data_instance)
  
      def generate_examples(self) -> Iterator[List[InputData]]:
          all_data: List[InputData] = []
          # Loading data from existing path if exists
          if self.config.output_path and os.path.exists(self.config.output_path):
              with open(self.config.output_path, 'rb') as file:
                  all_data = pickle.load(file)
                  for i in range(0, len(all_data), self.config.chunk_size):
                      yield all_data[i:i + self.config.chunk_size]
              return
  
          chunk: List[InputData] = []
  
          while len(all_data) < self.config.number_of_examples:
              messages = self.prepare_messages()
              message_batches = [
                  messages
              ] * (self.config.number_of_examples - len(all_data))
              with tqdm(
                  total=self.config.number_of_examples,
                  desc="Generating Examples",
                  unit="example"
              ) as pbar:
                  responses = asyncio.run(
                      utils.parallel_completions(
                          message_batches,
                          self.config.model_name,
                          self.config.max_token,
                          pbar=pbar
                      )
                  )
              for r in responses:
                  self.process_output(
                      r["choices"][0]["message"]["content"], all_data, chunk
                  )
              if chunk and len(chunk) >= self.config.chunk_size:
                  yield chunk
                  chunk = []
          if self.config.output_path:
              with open(self.config.output_path, 'wb') as file:
                  pickle.dump(all_data, file)
                  print(
                      f"Data succesfully generated and saved to {self.config.output_path}"
                  )
          if self.config.output_csv_path:
              with open(self.config.output_csv_path, 'w', newline='') as csvfile:
                  rows = [
                      BaseDataGenerator.input_data_to_csv_row(data)
                      for data in all_data
                  ]
                  header = rows[0].keys()
                  writer = csv.DictWriter(csvfile, fieldnames=header)
                  writer.writeheader()
                  for row in rows:
                      writer.writerow(row)
                  print(
                      f"Data succesfully generated and saved to {self.config.output_csv_path}"
                  )
          if chunk:
              yield chunk
  
  
  BaseDataGenerator.register_data_generator(
      "document_data_generator", DocumentDataGenerator,
      DocumentDataGeneratorConfig
  )
  
  
  def main():
      import time
      start_time = time.time()
      generator = DocumentDataGenerator(DocumentDataGenerator.default_config)
      res = generator.generate_examples()
  
      for d in res:
          print(d)
      end_time = time.time()
      elapsed_time = end_time - start_time
      print(f"Execution time: {elapsed_time:.2f} seconds")
  
  
  if __name__ == "__main__":
      main()
  
apis: langchain.text_splitter.RecursiveCharacterTextSplitter langchain.schema.Document langchain.document_loaders.GoogleDriveLoader langchain.document_loaders.UnstructuredFileLoader
extract_api: [((7802, 7926), 'yival.data_generators.base_data_generator.BaseDataGenerator.register_data_generator', 'BaseDataGenerator.register_data_generator', (['"""document_data_generator"""', 'DocumentDataGenerator', 'DocumentDataGeneratorConfig'], {}), "('document_data_generator',\n    DocumentDataGenerator, DocumentDataGeneratorConfig)\n", (7843, 7926), False, 'from yival.data_generators.base_data_generator import BaseDataGenerator\n'), ((1615, 2131), 'yival.schemas.data_generator_configs.DocumentDataGeneratorConfig', 'DocumentDataGeneratorConfig', ([], {'prompt': 'PROMPT_TEMPLATE', 'document': '""""""', 'source': '"""text"""', 'num_questions_per_chunk': '(5)', 'text_question_template': 'None', 'document_chunk_size': '(512)', 'number_of_examples': '(1)', 'question_gen_query': 'f"""You are a Teacher/Professor. Your task is to setup                         5 questions for an upcoming quiz/examination. The questions                         should be diverse in nature across the document. Restrict                         the questions to the context information provided."""'}), "(prompt=PROMPT_TEMPLATE, document='', source=\n    'text', num_questions_per_chunk=5, text_question_template=None,\n    document_chunk_size=512, number_of_examples=1, question_gen_query=\n    f'You are a Teacher/Professor. Your task is to setup                         5 questions for an upcoming quiz/examination. The questions                         should be diverse in nature across the document. Restrict                         the questions to the context information provided.'\n    )\n", (1642, 2131), False, 'from yival.schemas.data_generator_configs import DocumentDataGeneratorConfig\n'), ((7980, 7991), 'time.time', 'time.time', ([], {}), '()\n', (7989, 7991), False, 'import time\n'), ((8159, 8170), 'time.time', 'time.time', ([], {}), '()\n', (8168, 8170), False, 'import time\n'), ((3007, 3060), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': 'chunk_size'}), '(chunk_size=chunk_size)\n', (3037, 3060), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((2418, 2449), 'langchain.schema.Document', 'Document', ([], {'page_content': 'document'}), '(page_content=document)\n', (2426, 2449), False, 'from langchain.schema import Document\n'), ((5534, 5573), 'os.path.exists', 'os.path.exists', (['self.config.output_path'], {}), '(self.config.output_path)\n', (5548, 5573), False, 'import os\n'), ((2530, 2562), 'langchain.document_loaders.UnstructuredFileLoader', 'UnstructuredFileLoader', (['document'], {}), '(document)\n', (2552, 2562), False, 'from langchain.document_loaders import GoogleDriveLoader, UnstructuredFileLoader\n'), ((5664, 5681), 'pickle.load', 'pickle.load', (['file'], {}), '(file)\n', (5675, 5681), False, 'import pickle\n'), ((6126, 6217), 'tqdm.tqdm', 'tqdm', ([], {'total': 'self.config.number_of_examples', 'desc': '"""Generating Examples"""', 'unit': '"""example"""'}), "(total=self.config.number_of_examples, desc='Generating Examples', unit\n    ='example')\n", (6130, 6217), False, 'from tqdm import tqdm\n'), ((6977, 7004), 'pickle.dump', 'pickle.dump', (['all_data', 'file'], {}), '(all_data, file)\n', (6988, 7004), False, 'import pickle\n'), ((7470, 7512), 'csv.DictWriter', 'csv.DictWriter', (['csvfile'], {'fieldnames': 'header'}), '(csvfile, fieldnames=header)\n', (7484, 7512), False, 'import csv\n'), ((2709, 2747), 'langchain.document_loaders.GoogleDriveLoader', 'GoogleDriveLoader', ([], {'file_ids': '[document]'}), '(file_ids=[document])\n', (2726, 2747), False, 'from langchain.document_loaders import GoogleDriveLoader, UnstructuredFileLoader\n'), ((6345, 6451), 'yival.common.utils.parallel_completions', 'utils.parallel_completions', (['message_batches', 'self.config.model_name', 'self.config.max_token'], {'pbar': 'pbar'}), '(message_batches, self.config.model_name, self.\n    config.max_token, pbar=pbar)\n', (6371, 6451), False, 'from yival.common import utils\n'), ((7300, 7345), 'yival.data_generators.base_data_generator.BaseDataGenerator.input_data_to_csv_row', 'BaseDataGenerator.input_data_to_csv_row', (['data'], {}), '(data)\n', (7339, 7345), False, 'from yival.data_generators.base_data_generator import BaseDataGenerator\n')]