warm-start patch 2/3: find_positions.py

patched/vendor/ranking_moderation/scripts/find_positions.py

@@ -0,0 +1,482 @@
+import os
+import sys
+import logging
+from pathlib import Path
+import yaml
+import json
+import tqdm
+from logging.handlers import RotatingFileHandler
+from datetime import datetime
+import argparse
+from typing import Dict, Any, Set
+import copy
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# Add project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.append(str(project_root))
+
+# from src.ranking_moderation.bradley_terry_ranking import BradleyTerryRanking as Ranking
+from src.ranking_moderation.true_skill_ranking import TrueSkillRanking as Ranking
+from src.ranking_moderation.pairwise_comparison import PairwiseComparison, format_convs_as_xml, parse_unit_judgment, format_comments_as_xml, format_convs_as_xml_image
+from src.llm import AzureProxyOpenAI, LLMConfig
+
+logger = logging.getLogger(__name__)
+
+def setup_logging() -> logging.Logger:
+    """Configure logging with rotating file handlers."""
+    # Create logs directory if it doesn't exist
+    log_dir = Path("logs")
+    log_dir.mkdir(exist_ok=True)
+
+    # Configure logging
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S',
+        handlers=[
+            # Console handler - INFO level only
+            logging.StreamHandler(),
+            # Debug log file handler - DEBUG and above
+            RotatingFileHandler(
+                filename=log_dir / "debug.log",
+                maxBytes=10*1024*1024,  # 10MB
+                backupCount=5,
+                encoding='utf-8'
+            ),
+            # Info log file handler - INFO and above
+            RotatingFileHandler(
+                filename=log_dir / "info.log",
+                maxBytes=10*1024*1024,  # 10MB
+                backupCount=5,
+                encoding='utf-8'
+            ),
+            # Error log file handler - ERROR and above
+            RotatingFileHandler(
+                filename=log_dir / "error.log",
+                maxBytes=10*1024*1024,  # 10MB
+                backupCount=5,
+                encoding='utf-8'
+            )
+        ]
+    )
+
+    # Set console handler to INFO level
+    console_handler = logging.getLogger().handlers[0]
+    console_handler.setLevel(logging.ERROR)
+
+    # Set file handlers to their respective levels
+    debug_handler = logging.getLogger().handlers[1]
+    debug_handler.setLevel(logging.DEBUG)
+
+    info_handler = logging.getLogger().handlers[2]
+    info_handler.setLevel(logging.INFO)
+
+    error_handler = logging.getLogger().handlers[3]
+    error_handler.setLevel(logging.ERROR)
+
+    # Set all loggers to DEBUG level
+    for logger_name in logging.root.manager.loggerDict:
+        
+        logger = logging.getLogger(logger_name)
+        if logger_name in ("bytedes.mw", "elasticsearch", "urllib3.connectionpool"):
+            logger.setLevel(logging.ERROR)
+        else:
+            logger.setLevel(logging.DEBUG)
+        logger.propagate = True
+
+
+    return logging.getLogger(__name__)
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description='Find positions for new comments using Bradley-Terry ranking')
+    parser.add_argument('--config', type=str, required=True,
+                      help='Path to configuration YAML file')
+    parser.add_argument('--input_file', type=str, required=False, default=None,
+                      help='Input file (overrides yaml config field)')
+    parser.add_argument('--output_dir', type=str, required=False, default=None,
+                      help='output dir (overrides yaml config field)')
+    parser.add_argument('--pairwise_comparison', type=str, required=False, default=None,
+                      help='pairwise_comparison (overrides yaml config field)')
+    parser.add_argument('--ranking', type=str, required=False, default=None,
+                      help='ranking (overrides yaml config field)')
+    parser.add_argument('--es_index', type=str, required=False, default=None,
+                      help='es_index (overrides yaml config field)')
+    return parser.parse_args()
+
+def load_config(config_path: str) -> dict:
+    """Load configuration from YAML file."""
+    with open(config_path, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+def save_results(results: list, output_dir: Path, timestamp: str):
+    """Save results to both JSON and CSV formats."""
+    # Save JSON
+    json_path = output_dir / f"results_{timestamp}.json"
+    with json_path.open("w", encoding="utf-8") as f:
+        json.dump(results, f, ensure_ascii=False, indent=2)
+    logger.info(f"Saved results to {json_path}")
+
+def save_results_atomic(results, output_dir, timestamp):
+    # 最终文件
+    filename = output_dir / f"results_{timestamp}.json"
+    # 临时文件：只在 filename 后面加 ".tmp"
+    tmp_filename = filename.with_suffix(".json.tmp")
+
+    # 1) 写到临时文件
+    with open(tmp_filename, "w", encoding="utf-8") as f:
+        json.dump(results, f, ensure_ascii=False, indent=2)
+        f.flush()
+        os.fsync(f.fileno())
+
+    # 2) 原子替换
+    os.replace(tmp_filename, filename)
+
+def map_fields(instance: dict, field_mapping: dict) -> dict:
+    """Map fields according to configuration.
+    
+    Args:
+        instance: Original instance dictionary
+        field_mapping: Field mapping configuration
+        
+    Returns:
+        dict: Mapped instance with required fields
+    """
+    mapped_instance = {}
+    copy.deepcopy(instance)
+
+    # Map required fields
+    for target_field, source_field in field_mapping.items():
+        if isinstance(source_field, str):
+            if source_field in instance:
+                mapped_instance[target_field] = str(instance[source_field])
+            else:
+                logger.warning(f"Source field {source_field} not found in instance")
+                mapped_instance[target_field] = None
+        elif isinstance(source_field, list):
+            # Handle metadata fields
+            mapped_instance[target_field] = {}
+            for field in source_field:
+                if field in instance:
+                    mapped_instance[target_field][field] = instance[field]
+    
+    # Copy all original fields
+    instance.update(mapped_instance)
+    
+    return instance
+
+def load_test_items(input_file: Path, field_mapping: dict) -> list:
+    """Load test comments from JSON file and map fields.
+    
+    Args:
+        input_file: Path to the input JSON file
+        field_mapping: Field mapping configuration
+        
+    Returns:
+        List of mapped comment dictionaries
+    """
+    logger.info(f"Loading items from {input_file}")
+    items = []
+    try:
+        if input_file.suffix == ".json":
+            with input_file.open("r", encoding="utf-8") as f:                
+                data = json.load(f)
+                
+                # Handle different JSON structures
+                if isinstance(data, list):
+                    items = data
+                elif isinstance(data, dict) and "data" in data:
+                    items = data["data"]
+                else:
+                    raise ValueError(f"Unexpected JSON structure in {input_file}")                
+        elif input_file.suffix == ".jsonl":
+            try:
+                with input_file.open("r", encoding="utf-8") as f:
+                    for line in f:
+                        items.append(json.loads(line))
+            except json.JSONDecodeError:
+                # 如果按行解析失败，尝试一次性读取标准 JSON 格式
+                with input_file.open("r", encoding="utf-8") as f:
+                    items = json.load(f)
+        else:
+            raise ValueError(f"File {input_file} type is not recognized.")
+
+        # Map fields according to configuration
+        mapped_items = [map_fields(item, field_mapping) for item in items]
+        
+        logger.info(f"Loaded {len(mapped_items)} test items")
+        return mapped_items
+                
+    except Exception as e:
+        logger.error(f"Error loading {input_file}: {str(e)}")
+        raise
+
+def load_processed_items(output_dir: Path) -> Set[str]:
+    """Load previously processed comment IDs from output directory.
+    
+    Args:
+        output_dir: Path to the output directory containing processed results
+        
+    Returns:
+        Set of processed comment IDs
+    """
+    processed_ids = set()
+    
+    # Look for all JSON result files
+    result_files = list(output_dir.glob("results_*.json"))
+    for result_file in result_files:
+        try:
+            with result_file.open("r", encoding="utf-8") as f:
+                results = json.load(f)
+                for result in results:
+                    if "error" in result:
+                        continue
+                    if isinstance(result, dict) and "item" in result:
+                        item = result["item"]
+                        if isinstance(item, dict) and "item_id" in comment:
+                            processed_ids.add(str(comment[task_id]))
+        except Exception as e:
+            logger.error(f"Error loading processed comments from {result_file}: {str(e)}")
+            continue
+    
+    logger.info(f"Loaded {len(processed_ids)} previously processed comment IDs")
+    return processed_ids
+
+def initialize_llm(config: Dict[str, Any]) -> AzureProxyOpenAI:
+    """Initialize LLM with configuration.
+    
+    Args:
+        config: LLM configuration dictionary
+        
+    Returns:
+        AzureProxyOpenAI: Initialized LLM instance
+    """
+
+    llm_config = LLMConfig(
+        api_key=None,
+        api_base=config["api_url"],
+        model_name=config["model_name"],
+        caller=config["caller"],
+        max_tokens=config["max_tokens"],
+        temperature=config["temperature"]
+    )
+    return AzureProxyOpenAI(llm_config)
+
+
+def process_single_item(item, bt_ranker, config):
+    """
+    在线程中执行对单条评论的 find_sample_index 逻辑，并返回一个结果 dict。
+    如果发生异常，则返回包含错误信息的 dict。
+    """
+    # if not specified, by default use binary_search method
+    if "search_method" in config["ranking"]:
+        search_method = config["ranking"]["search_method"]
+    else:
+        search_method = 'similarity_search'
+
+    # warm-start：从 item 取 embedding 估出来的位置当 binary search 起点
+    start_pos = item.get("estimated_position")
+    if start_pos is not None:
+        try:
+            start_pos = int(start_pos)
+        except (TypeError, ValueError):
+            start_pos = None
+
+    try:
+        result = bt_ranker.find_sample_index(
+            new_item=item,
+            initial_comparisons=config["ranking"]["initial_comparisons"],
+            num_rounds=config["ranking"]["num_rounds"],
+            search_method=search_method,
+            start_pos=start_pos,
+        )
+        # 在外层不添加 timestamp 和 comment、processed_at，由主线程统一封装
+        return {
+            "success": True,
+            "payload": result
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e)
+        }
+
+def main():
+    # Parse command line arguments
+    args = parse_args()
+    
+    # Load configuration
+    config = load_config(args.config)
+    # config = load_config(project_root / "config" / "test_find_positions_config.yaml")
+    
+    # Override
+    # If CLI args provided override what's in YAML
+    if args.input_file:
+        if "data" in config:
+            config["data"]["input_file"] = args.input_file
+    if args.output_dir:
+        if "data" in config:
+            config["data"]["output_dir"] = args.output_dir
+    if args.pairwise_comparison:
+        if "data" in config and "storage" in config["data"]:
+            config["data"]["storage"]["pairwise_comparison"] = args.pairwise_comparison
+    if args.ranking:
+        if "data" in config and "storage" in config["data"]:
+            config["data"]["storage"]["ranking"] = args.ranking
+    if args.es_index:
+        if "data" in config and "storage" in config["data"]:
+            config["data"]["storage"]["es_index"] = args.es_index
+
+    
+
+    # Setup logging
+    global logger
+    logger = setup_logging()
+    
+    # Load prompt templates
+    prompt_path = project_root / config["data"]["prompt_template"]
+    with open(prompt_path, "r", encoding="utf-8") as f:
+        prompt_templates = yaml.safe_load(f)
+
+    # Initialize LLM
+    llm = initialize_llm(config["llm"])
+
+    # Setup storage directories
+    pairwise_storage_dir = project_root / config["data"]["storage"]["pairwise_comparison"]
+    ranking_storage_dir = project_root / config["data"]["storage"]["ranking"]
+    ranking_file_name = config["data"]["storage"].get("filename", "ruler_items.json")
+    output_dir = project_root / config["data"]["output_dir"]
+    pairwise_storage_dir.mkdir(parents=True, exist_ok=True)
+    ranking_storage_dir.mkdir(parents=True, exist_ok=True)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    business = config.get("business") or config.get("data", {}).get("business", "dm")
+
+    # Initialize PairwiseComparison
+    if business == "dm":
+        format_items = format_convs_as_xml
+    elif business == "dm_mm":
+        format_items = format_convs_as_xml_image
+    else:
+        format_items = format_comments_as_xml
+        
+    pairwise_comparison = PairwiseComparison(
+        llm=llm,
+        prompt_templates=prompt_templates,
+        format_items=format_items,
+        parse_judgment=parse_unit_judgment,
+        data_dir=str(pairwise_storage_dir),
+        es_psm=config["data"]["storage"]["es_psm"],
+        es_index=config["data"]["storage"]["es_index"],
+        local_cache_enabled=(config["data"]["storage"]["es_index"] in (None, "None")),
+        max_comparisons_per_pair=config["ranking"]["max_comparisons_per_pair"],
+        max_workers=config["ranking"]["max_workers"],
+        max_retries=config["ranking"]["max_retries"],
+        initial_retry_delay=config["ranking"]["initial_retry_delay"],
+        business=business
+    )
+
+    # Initialize BradleyTerryRanking
+    bt_ranker = Ranking(
+        pairwise_comparison=pairwise_comparison,
+        data_dir=str(ranking_storage_dir),
+        file_name=str(ranking_file_name),
+        convergence_threshold=config["ranking"]["convergence_threshold"]
+    )
+    
+    # Load previously processed comments
+    processed_ids = load_processed_items(output_dir)
+    
+    # Load test comments with field mapping
+    input_file = project_root / config["data"]["input_file"]
+    test_items = load_test_items(input_file, config["data"]["field_mapping"])
+    # random.shuffle(test_comments)
+    # test_comments = test_comments[:5]
+
+    if not test_items:
+        logger.error("No test items loaded. Exiting.")
+        return
+    
+    # Filter out already processed comments
+    test_items = [
+        item for item in test_items 
+        if str(item["item_id"]) not in processed_ids
+    ]
+    
+    if not test_items:
+        logger.info("All items have been processed. Exiting.")
+        return
+    
+    logger.info(f"Found {len(test_items)} new comments to process")
+    
+    # Initialize results list
+    results = []
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    future_to_item = {}
+
+    with ThreadPoolExecutor(max_workers=config["ranking"]["max_workers"]) as executor:
+        # 2) 为每条评论提交一个任务
+        for item in test_items:
+            future = executor.submit(process_single_item, item, bt_ranker, config)
+            future_to_item[future] = item
+
+        processed_count = 0
+        # 3) as_completed 会在每个 future 运行完毕时返回
+        for future in tqdm.tqdm(as_completed(future_to_item),
+                                 total=len(test_items),
+                                 desc="Processing comments"):
+            comment = future_to_item[future]
+            # 4) 在这里拿到单条评论的结果（或异常）
+            try:
+                outcome = future.result()
+            except Exception as e:
+                # 如果连 process_single_comment 本身抛出了未捕获异常，也把它当作一个 error 结果处理
+                outcome = {
+                    "success": False,
+                    "error": f"Uncaught exception: {str(e)}"
+                }
+            processed_count += 1
+            # 5) 根据 outcome 构造要保存到 results 列表里的条目
+            if outcome["success"]:
+                result = outcome["payload"]
+                # 把原始评论内容、处理时间戳补回去
+                result["comment"] = comment
+                result["processed_at"] = datetime.now().isoformat()
+                results.append(result)
+
+                # 记录日志
+                pos = result.get("position", None)
+                score = result.get("score", None)
+                rounds_done = result.get("rounds_completed", None)
+                if pos is not None:
+                    logger.debug(f"Found position for item {item.get('item_id', '')}: {pos}")
+                if score is not None:
+                    logger.debug(f"Final score: {score:.4f}")
+                if rounds_done is not None:
+                    logger.debug(f"Rounds completed: {rounds_done}")
+            else:
+                # 出现了业务层面的错误
+                error_entry = {
+                    "item": item,
+                    "error": outcome.get("error"),
+                    "processed_at": datetime.now().isoformat()
+                }
+                results.append(error_entry)
+                logger.error(f"Error processing item {item.get('item_id', '')}: {outcome.get('error')}")
+            
+            # 6) 每处理完一条，就把当前 results 列表写盘
+            if processed_count % 10 ==0:
+                save_results_atomic(results, output_dir, timestamp)
+        
+    save_results_atomic(results, output_dir, timestamp)
+    logger.info(f"Completed processing {len(test_items)} items")
+    logger.info(f"Results saved to {output_dir}")
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        logger.exception("An error occurred during execution")
+        raise