"""
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()

langchain.document_loaders.GoogleDriveLoader langchain.schema.Document langchain.document_loaders.UnstructuredFileLoader langchain.text_splitter.RecursiveCharacterTextSplitter
[((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')]