src/components/data_extractor.py - master - Walmet/CTF

src/components/data_extractor.py
167 строк6 KB
Maksim Goncharov
MVP
14 дек 2025, 23:49
14 дек 2025, 23:49a4c2b1d
import pandas as pd
import numpy as np
from typing import Tuple, Dict, Any, Union
import random

from .data_validator import DataValidator
from ..models.model import QAModel

class DataExtractor:
    """
    Class for extracting, transforming, and creating metadata for documents.
    """

    def __init__(self, default_template = None):
        self.default_template = default_template
        self.data_validator = DataValidator()
        self.qa_model = QAModel()

    def _generate_mock_metadata(self, df: pd.DataFrame, file_name: str) -> pd.DataFrame:
        """
        A helper function that simulates the neural network logic by creating
        MultiIndex metadata (cell-by-cell status).

        Status:
        0: No source link (data comes from a template). Background: Yellow.
        1: Recognition successful (confident). Background: White.
        2: Recognition unsuccessful (low confidence, requires manual editing). Background: Lavender/Purple.
        """
        df_transformed = df.astype(str).reset_index(drop=True)
        num_rows = df_transformed.shape[0]

        index_tuples = []
        source_data = {
            'source_file': [],
            'source_row': [],
            'source_col': [],
            'status': []
        }

        for row_idx in range(num_rows):
            for col_idx in range(len(df_transformed.columns)):
                index_tuples.append((row_idx, col_idx))

                source_data['source_file'].append(file_name)
                source_data['source_row'].append(row_idx)
                source_data['source_col'].append(col_idx)
                source_data['status'].append(np.random.choice([0, 1, 2]))

        multi_index = pd.MultiIndex.from_tuples(index_tuples, names=['Row', 'Col'])

        metadata_df = pd.DataFrame({
            'source_file': source_data['source_file'],
            'source_row': source_data['source_row'],
            'source_col': source_data['source_col'],
            'status': source_data['status']
        }, index=multi_index)

        return metadata_df

    def predict(self, text, prompt):
        generated, probability = self.qa_model.predict(text, prompt)
        return generated, probability

    def generate(self, dataframes, prompt):
        max_generated = 0
        max_probability = 0
        max_source_file = ''
        max_source_row = 0
        max_source_col = ''

        for source_file in dataframes:
            dataframe = dataframes[source_file].astype(str)
            for row in range(dataframe.shape[0]):
                for col in range(dataframe.shape[1]):
                    text = dataframe.iloc[row, col]
                    generated, probability = self.predict(text, prompt)
                    if probability > max_probability:
                        max_generated = generated
                        max_probability = probability
                        max_source_file = source_file
                        max_source_row = row
                        max_source_col = col

        return max_generated, max_probability, max_source_file, max_source_row, max_source_col


    def get_prompt(self, text):
        prompt_start = "#PROMPT# "
        if text.startswith(prompt_start):
            return str(text[len(prompt_start):])
        return None

    def predict_and_generate_metadata(self, dataframes, transformed_dataframe):
        index_tuples = []
        source_data = {
            'source_file': [],
            'source_row': [],
            'source_col': [],
            'status': []
        }

        for row in range(transformed_dataframe.shape[0]):
            for col in range(transformed_dataframe.shape[1]):
                index_tuples.append((row, col))

                text = str(transformed_dataframe.iloc[row, col])
                prompt = self.get_prompt(text)
                if prompt:
                    generated, probability, source_file, source_row, source_col = self.generate(dataframes, prompt)

                    if not self.data_validator.verify(generated, prompt, probability):
                        source_data['status'].append(2)
                        source_data['source_file'].append('')
                        source_data['source_row'].append(0)
                        source_data['source_col'].append('')
                        continue

                    if self.data_validator.validate(generated, prompt, probability):
                        source_data['status'].append(1)
                    else:
                        source_data['status'].append(2)
                    source_data['source_file'].append(source_file)
                    source_data['source_row'].append(source_row)
                    source_data['source_col'].append(source_col)

                    transformed_dataframe.iloc[row, col] = generated
                else:
                    source_data['status'].append(0)
                    source_data['source_file'].append('')
                    source_data['source_row'].append(0)
                    source_data['source_col'].append('')

        multi_index = pd.MultiIndex.from_tuples(index_tuples, names=['Row', 'Col'])

        metadata_df = pd.DataFrame({
            'source_file': source_data['source_file'],
            'source_row': source_data['source_row'],
            'source_col': source_data['source_col'],
            'status': source_data['status']
        }, index=multi_index)

        return metadata_df


    def transform_dataframes(self, dataframes: Dict[str, pd.DataFrame], template: str = None) -> Tuple[
        Union[pd.DataFrame, None], Union[pd.DataFrame, None]]:
        """
        Accepts a set of documents and simulates their transformation and validation.
        """
        if template is None and self.default_template is not None:
            template = self.default_template

        if not isinstance(dataframes, dict) or not dataframes or template is None:
            return None, None

        # file_name, df = next(iter(dataframes.items()))
        # transformed_dataframe = df.astype(str).reset_index(drop=True)
        # metadata_df = self._generate_mock_metadata(transformed_dataframe, file_name)

        try:
            transformed_dataframe = pd.read_csv(template, index_col=None, dtype=str)
        except Exception as e:
            return None, None

        metadata_df = self.predict_and_generate_metadata(dataframes, transformed_dataframe)

        return transformed_dataframe, metadata_df