Python LLM task classifier V1

#!/usr/bin/env python
"""
Production Training Data Preparation Script (Revised)

Features Included:
  - Granular chunking of text, with optional summarization for large paragraphs.
  - Lightweight task quality scoring & filtering that won't significantly slow down processing.
  - Simple keyword-based domain classification for each task.
  - Enhanced logging with rotating file handler & incremental writes (partial Parquet).
  - Bias detection with rebalancing feedback & early exits if triggered too often.
  - Parallel file processing with resource safeguards.

Excluded:
  - Interactive Preview & Confirmation (explicitly not desired).
"""

import os
import glob
import argparse
import random
import sys
import logging
import subprocess
import importlib

# Function to install missing packages
def install_package(package_name):
    print(f"Attempting to install {package_name}...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
        print(f"Successfully installed {package_name}")
        return True
    except subprocess.CalledProcessError:
        print(f"Failed to install {package_name}. Please install it manually.")
        return False

# Dictionary mapping module names to their pip package names
package_map = {
    "yaml": "PyYAML",
    "pandas": "pandas",
    "torch": "torch",
    "transformers": "transformers",
    "PyPDF2": "PyPDF2",
    "docx": "python-docx",
    "tqdm": "tqdm",
    "psutil": "psutil",
}

# Try importing each package, install if missing
for module_name, package_name in package_map.items():
    try:
        globals()[module_name] = importlib.import_module(module_name)
    except ImportError:
        print(f"{module_name} not found. Installing {package_name}...")
        if install_package(package_name):
            try:
                globals()[module_name] = importlib.import_module(module_name)
                print(f"Successfully imported {module_name} after installation")
            except ImportError:
                print(f"Failed to import {module_name} even after installing {package_name}. Exiting.")
                sys.exit(1)
        else:
            print(f"Could not install {package_name}. Exiting.")
            sys.exit(1)

# Now that we have all imports, import specific items from modules
from logging.handlers import RotatingFileHandler
from concurrent.futures import ThreadPoolExecutor, as_completed
from transformers import AutoTokenizer, AutoModel, pipeline

# Define tqdm fallback if still not available
if 'tqdm' in globals():
    from tqdm import tqdm
else:
    tqdm = lambda x, **kwargs: x  # Fallback if tqdm is not installed

import psutil  # for resource safeguard checks

#########################################
# Configuration, Logging & Utility Functions
#########################################

def load_config(config_path):
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config

def setup_logging(logging_config):
    log_level = getattr(logging, logging_config.get("level", "INFO").upper(), logging.INFO)
    logger = logging.getLogger()
    logger.setLevel(log_level)
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")

    # File handler with rotation if enabled
    if logging_config.get("log_file"):
        log_file = logging_config["log_file"]
        if logging_config.get("rotate", False):
            max_bytes = parse_size(logging_config.get("max_file_size", "100MB"))
            backup_count = logging_config.get("backup_count", 5)
            file_handler = RotatingFileHandler(log_file, maxBytes=max_bytes, backupCount=backup_count)
        else:
            file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    
    if logging_config.get("enable_console", True):
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)

def parse_size(size_str):
    size_str = size_str.strip().upper()
    if size_str.endswith("GB"):
        return int(float(size_str[:-2]) * 1024**3)
    elif size_str.endswith("MB"):
        return int(float(size_str[:-2]) * 1024**2)
    elif size_str.endswith("KB"):
        return int(float(size_str[:-2]) * 1024)
    else:
        return int(size_str)

#########################################
# Text Extraction & Summarization
#########################################

def extract_text(file_path):
    """Extract text from a file (PDF, DOCX, or TXT)."""
    ext = os.path.splitext(file_path)[1].lower()
    if ext == ".txt":
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            logging.error(f"Error reading {file_path}: {e}")
            return ""
    elif ext == ".pdf":
        text = ""
        try:
            with open(file_path, 'rb') as f:
                reader = PyPDF2.PdfReader(f)
                for page in reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        text += page_text + "\n"
        except Exception as e:
            logging.error(f"Error reading PDF {file_path}: {e}")
        return text
    elif ext == ".docx":
        try:
            doc = docx.Document(file_path)
            return "\n".join([p.text for p in doc.paragraphs])
        except Exception as e:
            logging.error(f"Error reading DOCX {file_path}: {e}")
            return ""
    else:
        logging.warning(f"Unsupported file type: {file_path}")
        return ""

def clean_and_chunk_text(text, max_chunk_chars=500):
    """
    Clean text and split it into smaller chunks. If summarization is enabled,
    we optionally summarize any chunk beyond a threshold.
    """
    chunks = []
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    for para in paragraphs:
        if len(para) > max_chunk_chars:
            # We'll further split or optionally summarize
            sentences = para.split(". ")
            current_chunk = ""
            for sentence in sentences:
                if len(current_chunk) + len(sentence) + 2 <= max_chunk_chars:
                    current_chunk += sentence + ". "
                else:
                    chunks.append(current_chunk.strip())
                    current_chunk = sentence + ". "
            if current_chunk:
                chunks.append(current_chunk.strip())
        else:
            chunks.append(para)
    return chunks

def optional_summarize_chunk(chunk, summarizer, max_length=130, min_length=30):
    """
    If summarizer is enabled in config, we run a summarization step
    on the chunk if chunk is beyond a certain length.
    """
    if len(chunk.split()) > 512:  # arbitrary threshold
        try:
            summary = summarizer(chunk, max_length=max_length, min_length=min_length, do_sample=False)
            return summary[0]['summary_text']
        except Exception as e:
            logging.error(f"Summarization failed: {e}")
            return chunk
    else:
        return chunk

#########################################
# Domain Classification (Keyword-based)
#########################################

def classify_domain(task_description, domain_keywords):
    """
    Simple approach: check for keywords in the description, pick whichever domain
    has the highest count of matches. If none match, default to "General".
    """
    text_lower = task_description.lower()
    best_domain = "General"
    best_count = 0
    for domain, keywords in domain_keywords.items():
        count = sum(text_lower.count(k.lower()) for k in keywords)
        if count > best_count:
            best_count = count
            best_domain = domain
    return best_domain

#########################################
# Quality Scoring & Filtering
#########################################

def score_task_quality(task):
    """
    Minimal overhead scoring:
      - If description is too short, penalty
      - If repeated text is found, penalty
    We'll produce a float in [0..1].
    """
    desc = task["task_description"].strip()
    if len(desc) < 20:
        return 0.2  # too short
    # simple repeated text check
    tokens = desc.split()
    if len(tokens) > 10:
        unique_tokens = set(tokens)
        rep_ratio = (len(tokens) - len(unique_tokens)) / float(len(tokens))
        # if more than 30% tokens repeated, penalty
        if rep_ratio > 0.3:
            return 0.3
    return 0.8  # default pass

#########################################
# Task Generation & Rebalancing
#########################################

def generate_tasks_for_subject(subject, content, config, domain_keywords):
    tasks = []
    domain_distribution = {}
    rebalancing_count = 0

    batch_size = config["task_generation"].get("batch_size", 100)
    min_tasks = config["task_generation"].get("min_tasks_per_subject", 50000)
    max_tasks = config["task_generation"].get("max_tasks_per_subject", 100000)
    bias_threshold = config["bias_detection"].get("threshold", 0.1)
    min_tasks_before_rebalancing = config["bias_detection"].get("min_tasks_before_rebalancing", 5000)
    max_rebalancing_attempts = config["bias_detection"].get("max_rebalancing_attempts", 5)
    user_confirmation = config["bias_detection"].get("user_confirmation", False)

    quality_threshold = config["task_generation"].get("quality_threshold", 0.0)

    # chunk & optionally summarize
    chunks = clean_and_chunk_text(content, max_chunk_chars=500)
    summarizer_enabled = config.get("summarization", {}).get("enabled", False)

    # if summarizer is enabled, initialize pipeline once
    summarizer_pipeline = None
    if summarizer_enabled:
        model_name = config["summarization"].get("model_name", "facebook/bart-large-cnn")
        try:
            summarizer_pipeline = pipeline("summarization", model=model_name)
        except Exception as e:
            logging.error(f"Failed to load summarizer model: {e}")
            summarizer_enabled = False

    # For demonstration, we create tasks from each chunk + repeated until min_tasks
    total_tasks = 0
    iteration_count = 0

    while total_tasks < min_tasks:
        batch_tasks = []
        for _ in range(batch_size):
            if not chunks:
                # no content, break
                break
            chunk = random.choice(chunks)
            if summarizer_enabled:
                chunk = optional_summarize_chunk(chunk, summarizer_pipeline)
            # minimal "task" generation
            task_description = f"Review content snippet: '{chunk[:50]}...'"
            # domain classification
            domain = classify_domain(task_description, domain_keywords)
            recommended_execution = "automated analysis" if len(chunk) > 100 else "manual review"

            task = {
                "task_description": task_description,
                "domain": domain,
                "recommended_execution": recommended_execution,
            }
            # Score & filter
            quality_score = score_task_quality(task)
            if quality_score >= quality_threshold:
                batch_tasks.append(task)
                domain_distribution[domain] = domain_distribution.get(domain, 0) + 1

        tasks.extend(batch_tasks)
        total_tasks = len(tasks)

        # domain bias detection
        expected_share = total_tasks / max(1, len(domain_keywords.keys()) + 1)
        # +1 to avoid zero if no domain keywords
        for d, count in domain_distribution.items():
            if count > expected_share * (1 + bias_threshold):
                allowed = int(expected_share * (1 + bias_threshold))
                # filter tasks to remove extra occurrences6
                rebalanced_tasks = []
                current_count = 0
                for t in tasks:
                    if t["domain"] == d:
                        if current_count < allowed:
                            rebalanced_tasks.append(t)
                            current_count += 1
                    else:
                        rebalanced_tasks.append(t)
                        
                # Replace the tasks list with the rebalanced list
                tasks = rebalanced_tasks
                total_tasks = len(tasks)
                
                # Recalculate the entire domain distribution after rebalancing
                domain_distribution = {}
                for task in tasks:
                    domain = task["domain"]
                    domain_distribution[domain] = domain_distribution.get(domain, 0) + 1
                
                rebalancing_count += 1
                logging.info(f"Bias detected in domain '{d}'. Rebalancing... (Attempt: {rebalancing_count})")
                break

        if rebalancing_count > max_rebalancing_attempts and total_tasks < min_tasks_before_rebalancing:
            logging.error("Exceeded maximum rebalancing attempts before reaching minimum tasks. Exiting early.")
            sys.exit(1)

        iteration_count += 1
        # Resource safeguard check
        check_resource_usage(config)

        # break if we are out of content or reached max
        if total_tasks >= max_tasks:
            tasks = tasks[:max_tasks]
            break
        if not chunks:
            break

    return tasks

def check_resource_usage(config):
    """
    If parallel_processing.resource_limits are set, check if memory or CPU usage
    exceeds thresholds. If so, exit or raise an error.
    """
    resource_limits = config.get("parallel_processing", {}).get("resource_limits", {})
    if not resource_limits:
        return

    max_memory_mb = resource_limits.get("max_memory_mb", 0)
    max_cpu_percent = resource_limits.get("max_cpu_percent", 0)

    if max_memory_mb > 0:
        mem_info = psutil.virtual_memory()
        used_mb = (mem_info.total - mem_info.available) / (1024 * 1024)
        if used_mb > max_memory_mb:
            logging.error(f"Memory usage exceeded {max_memory_mb} MB limit. Current usage: {used_mb:.2f} MB.")
            sys.exit(1)

    if max_cpu_percent > 0:
        cpu_percent = psutil.cpu_percent(interval=0.1)
        if cpu_percent > max_cpu_percent:
            logging.error(f"CPU usage exceeded {max_cpu_percent}% limit. Current usage: {cpu_percent:.2f}%.")
            sys.exit(1)

#########################################
# Append Tasks to Parquet (Incremental Writes)
#########################################

def append_tasks_to_parquet(subject, new_tasks, output_folder, file_count):
    df_new = pd.DataFrame(new_tasks)
    output_path = os.path.join(output_folder, f"{subject}_{file_count}.parquet")
    df_new.to_parquet(output_path, index=False)
    return output_path

#########################################
# Main Function
#########################################

def main():
    parser = argparse.ArgumentParser(description="Production Training Data Preparation Script (Revised)")
    parser.add_argument("--config", type=str, default="production_config.yaml", help="Path to YAML config file")
    parser.add_argument("--no_fp16", action="store_true", help="Disable FP16 training")
    args = parser.parse_args()

    # Load config and setup logging
    config = load_config(args.config)
    setup_logging(config.get("logging", {}))
    logging.info("Configuration loaded. Logging set up.")

    # Some device config if needed (we keep minimal usage)
    device_config = config.get("device_config", {})
    # we won't do advanced model usage here unless summarizer is enabled

    input_folder = config.get("storage", {}).get("input_folder", "")
    if not os.path.isdir(input_folder):
        logging.error("The provided folder path is not valid.")
        sys.exit(1)

    output_folder = config.get("storage", {}).get("output_folder", os.path.join(input_folder, "parquet_output"))
    os.makedirs(output_folder, exist_ok=True)

    file_extensions = ("*.txt", "*.pdf", "*.docx")
    file_paths = []
    for ext in file_extensions:
        file_paths.extend(glob.glob(os.path.join(input_folder, ext)))
    if not file_paths:
        logging.error("No supported document files found in the folder.")
        sys.exit(1)

    domain_keywords = config.get("subject_extraction", {}).get("keywords", {})

    # parallel processing
    parallel_config = config.get("parallel_processing", {})
    max_workers = parallel_config.get("num_workers", 4)
    enabled_parallel = parallel_config.get("enabled", True)

    subjects_dict = {}
    if enabled_parallel:
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_file = {executor.submit(extract_text, file): file for file in file_paths}
            for future in tqdm(as_completed(future_to_file), total=len(file_paths), desc="Ingesting files"):
                file = future_to_file[future]
                text = future.result()
                if not text.strip():
                    continue
                # for simplicity, just classify entire doc as "General" subject
                # or use filename as subject
                subject_name = os.path.splitext(os.path.basename(file))[0]
                if subject_name in subjects_dict:
                    subjects_dict[subject_name] += "\n" + text
                else:
                    subjects_dict[subject_name] = text
                check_resource_usage(config)
    else:
        # no parallel
        for file in tqdm(file_paths, desc="Ingesting files"):
            text = extract_text(file)
            if not text.strip():
                continue
            subject_name = os.path.splitext(os.path.basename(file))[0]
            if subject_name in subjects_dict:
                subjects_dict[subject_name] += "\n" + text
            else:
                subjects_dict[subject_name] = text
            check_resource_usage(config)

    # Now generate tasks for each subject
    file_count = 1
    for subject, content in subjects_dict.items():
        logging.info(f"Generating tasks for subject: {subject}")
        new_tasks = generate_tasks_for_subject(subject, content, config, domain_keywords)
        # We do an incremental write for each subject
        output_path = append_tasks_to_parquet(subject, new_tasks, output_folder, file_count)
        logging.info(f"Saved {len(new_tasks)} tasks for subject '{subject}' at {output_path}")
        file_count += 1

if __name__ == "__main__":
    main()