bm25.py

from rank_bm25 import BM25Okapi
from multiprocessing import Pool
from collections import defaultdict
from typing import List
from datasets import CodeBlock
from functools import partial
import math

import logging
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

def split_into_smaller_blocks(code_block, enable_fixed_block):
    """
    Split large blocks of code into smaller ones, each containing no more than 12 non-empty lines.
    """
    smaller_blocks = []

    if enable_fixed_block:
        lines = [line for line in code_block.code_content.split('\n') if line.strip() != '']
        for i in range(0, min(len(lines),5000), 12):
            start_line_offset = i
            end_line_offset = min(i + 12, len(lines))
            block_content = '\n'.join(lines[start_line_offset:end_line_offset])
            smaller_blocks.append(CodeBlock(code_block.file_path, 
                                            f"file path: {code_block.file_path}\nlines: {start_line_offset}-{end_line_offset - 1}",
                                            block_content,
                                            code_block.language,
                                            'fixed_block'))

    else:
        # Split the code by spaces, then reassemble it into blocks.
        mini_blocks = []
        current_block = [] 
        for line in code_block.code_content.splitlines(): 
            if line.strip() == '':  
                if current_block: 
                    mini_blocks.append(current_block)
                    current_block = []
            else:
                current_block.append(line)
        if current_block: 
            mini_blocks.append(current_block)

        max_len = 15
        temp_mini_blocks = []
        for mini_block in mini_blocks:
            if len(mini_block) > max_len:
                for idx in range(0, len(mini_block), max_len):
                    temp_mini_blocks.append(mini_block[idx: idx+max_len])
            else:
                temp_mini_blocks.append(mini_block)
        mini_blocks = temp_mini_blocks

        current_content = []
        total_lines = 0  
        for block in mini_blocks:
            if total_lines >= 5000:  
                break  
            if len(current_content) + len(block) <= 15:  
                current_content.extend(block)
                total_lines += len(block)  
            else:  
                if current_content:  
                    smaller_blocks.append(CodeBlock(code_block.file_path, 
                                                    f"file path: {code_block.file_path}\nlines: {total_lines-len(current_content)+1}-{total_lines}",
                                                    '\n'.join(current_content),
                                                    code_block.language,
                                                    'mini_block'))
                current_content = block  
                total_lines += len(block)  
        if current_content:  
            smaller_blocks.append(CodeBlock(code_block.file_path, 
                                            f"file path: {code_block.file_path}\nlines: {total_lines-len(current_content)+1}-{total_lines}",
                                            '\n'.join(current_content),
                                            code_block.language,
                                            'mini_block'))
        
    return smaller_blocks

class TaskSpecificBM25:
    def __init__(self, examples, args):
        self.bm25_indices = {}
        self.code_blocks = {}
        self._build_indices(examples, args)
        self.args = args
        
    def _build_indices(self, examples, args):
        num_processes = 32  
        num_examples_per_batch = math.ceil(len(examples) / num_processes)

    
        example_batches = [examples[i:i + num_examples_per_batch] for i in range(0, len(examples), num_examples_per_batch)]

        with Pool(processes=num_processes) as pool:
            # results = pool.map(self._process_batch, example_batches, args.enable_fixed_block)
            # results = pool.map(lambda batch: self._process_batch(batch, args.enable_fixed_block), example_batches)
            from functools import partial
            partial_process_batch = partial(self._process_batch, enable_fixed_block=args.enable_fixed_block)
            results = pool.map(partial_process_batch, example_batches)
        
        for batch_result in results:
            for task_id, code_blocks, bm25_index in batch_result:
                self.bm25_indices[task_id] = bm25_index
                self.code_blocks[task_id] = code_blocks

        block_len = 0
        block_num = 0
        for task_id, code_blocks in self.code_blocks.items():
            block_num += len(code_blocks)
            for code_block in code_blocks:
                block_len += len(code_block.code_content.splitlines())
        
        logger.info(f'Block avg line: {round(block_len / block_num, 2)}')

    @staticmethod
    def _process_batch(batch, enable_fixed_block):
        batch_result = []
        for example in batch:
            code_blocks = []
            for code_block in example.related_files:
                code_blocks.extend(split_into_smaller_blocks(code_block, enable_fixed_block))
            
            bm25_index = None
            if len(code_blocks) != 0:
                bm25_index = BM25Okapi([code_block.code_content.lower().split() for code_block in code_blocks])
            
            batch_result.append((example.task_id, code_blocks, bm25_index))
        return batch_result

    def query(self, task_ids: List[int], keyword_lists: List[List[str]], topk: int):
        """
        Use key word lists to retrieve. 
        
        Args:
            task_ids: task id lists
            keyword_lists: key word lists(only a sample)
            topk: number of best top k retrieval results
        
        Returns:
            top_k lists
        """
        results = []
        for task_id, keywords in zip(task_ids, keyword_lists):
            bm25_index = self.bm25_indices.get(task_id)
            if bm25_index:
                # 直接使用关键词列表，不需要分词
                scores = bm25_index.get_scores(keywords)

                sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
                fixed_block_indices = []
                signature_block_indices = []
                full_func_block_indices = []
                mini_block_indices = []
                for i in sorted_indices:
                    block = self.code_blocks[task_id][i]
                    if block._type == "fixed_block":
                        fixed_block_indices.append(i)
                    elif block._type == "signature_block":
                        signature_block_indices.append(i)
                    elif block._type == "full_func_block":
                        full_func_block_indices.append(i)
                    elif block._type == "mini_block":
                        mini_block_indices.append(i)
                    
                topk_fixed_blocks = [self.code_blocks[task_id][i] for i in fixed_block_indices[:topk]]
                topk_mini_blocks = [self.code_blocks[task_id][i] for i in mini_block_indices[:topk]]

                if self.args.enable_fixed_block:
                    task_results = topk_fixed_blocks
                else:
                    task_results = topk_mini_blocks
                results.append(task_results)
            else:
                results.append([])
        return results

    def query_batch(self, task_ids: List[int], queries_list: List[List[List[str]]], topk: int):
        """
        Like query, but support multi-sampling
        
        :param task_ids: task_id [num_examples]
        :param queries_list: shape is [num_examples][num_samples][keywords_per_sample]
        :return: shape is [num_examples][topk]
        """
        results = []
        for task_id, sample_queries in zip(task_ids, queries_list):
            sample_results = []
            
            # 对每个采样关键词单独检索
            for query in sample_queries:
                bm25_index = self.bm25_indices.get(task_id)
                if not bm25_index:
                    sample_results.append([])
                    continue
                
                # 执行BM25查询，将关键词列表作为一个整体查询
                scores = bm25_index.get_scores(query)
                
                # 按分数排序
                sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
                
                # 应用块类型过滤策略
                fixed_blocks = []
                mini_blocks = []
                for i in sorted_indices:
                    block = self.code_blocks[task_id][i]
                    if block._type == "fixed_block":
                        fixed_blocks.append(block)
                    elif block._type == "mini_block":
                        mini_blocks.append(block)
                
                # 根据参数选择块类型
                if self.args.enable_fixed_block:
                    final_blocks = fixed_blocks[:topk]
                else:
                    final_blocks = mini_blocks[:topk]
                
                sample_results.append(final_blocks)
            
            results.append(sample_results)

        return results


    # def query(self, task_ids: List[int], queries: List[str], topk: int):
    #     results = []
    #     for task_id, query in zip(task_ids, queries):
    #         bm25_index = self.bm25_indices.get(task_id)
    #         if bm25_index:
    #             query_tokens = query.split()
    #             scores = bm25_index.get_scores(query_tokens)

    #             # topk_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:topk]
    #             # task_results = [self.code_blocks[task_id][i] for i in topk_indices]
    #             # results.append(task_results)

    #             #  ||
    #             # \||/
    #             #  \/

    #             sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    #             fixed_block_indices = []
    #             signature_block_indices = []
    #             full_func_block_indices = []
    #             mini_block_indices = []
    #             for i in sorted_indices:
    #                 block = self.code_blocks[task_id][i]
    #                 if block._type == "fixed_block":
    #                     fixed_block_indices.append(i)
    #                 elif block._type == "signature_block":
    #                     signature_block_indices.append(i)
    #                 elif block._type == "full_func_block":
    #                     full_func_block_indices.append(i)
    #                 elif block._type == "mini_block":
    #                     mini_block_indices.append(i)
                    
    #             topk_fixed_blocks = [self.code_blocks[task_id][i] for i in fixed_block_indices[:topk]]
    #             # topk_signature_blocks = [self.code_blocks[task_id][i] for i in signature_block_indices[:topk]]
    #             # topk_full_func_blocks = [self.code_blocks[task_id][i] for i in full_func_block_indices[:topk]]
    #             topk_mini_blocks = [self.code_blocks[task_id][i] for i in mini_block_indices[:topk]]

    #             if self.args.enable_fixed_block:
    #                 task_results = topk_fixed_blocks
    #             else:
    #                 task_results = topk_mini_blocks
    #             results.append(task_results)
    #         else:
    #             results.append([])
    #         #results.append([CodeBlock("__init__.py", 0, 0, "") for _ in range(topk+1)])
    #     return results