Models API¶

@staticmethod
@handle_exceptions(log_exceptions=True, include_exc_info=True)
def list_available_models() -> Mapping[str, Dict[str, str]]:
    """
    List available pre-trained cross-encoder models categorized by type.

    This static method returns a dictionary of model categories and their
    corresponding model recommendations that can be used with the search system.
    These models are defined in the ALTERNATIVE_MODELS constant.

    Returns:
        Mapping[str, Dict[str, str]]: A dictionary where:
            - Keys are model categories (e.g., "Semantic Search", "Question Answering")
            - Values are dictionaries mapping model names to descriptions

    Example:
        ```python
        # Get a dictionary of available models by category
        available_models = BaseSearchModel.list_available_models()

        # Print model categories and models
        for category, models in available_models.items():
            print(f"\n{category}:")
            for model_name, description in models.items():
                print(f"  - {model_name}: {description}")
        ```
    """
    return ALTERNATIVE_MODELS

@staticmethod
@handle_exceptions(log_exceptions=True, include_exc_info=True)
def list_fine_tuned_models() -> Dict[str, str]:
    """
    List locally available fine-tuned models that can be used for search.

    This static method scans the fine-tuned model directory to find models
    that have been fine-tuned specifically for anime/manga search. It identifies
    valid models by checking for the presence of a config.json file.

    Returns:
        Dict[str, str]: A dictionary mapping:
            - Keys: Model directory names
            - Values: Full paths to the model directories

    Notes:
        - Searches in the "model/fine-tuned" directory by default
        - Only directories containing a config.json file are included
        - Returns an empty dictionary if no fine-tuned models are found

    Example:
        ```python
        # Get a dictionary of available fine-tuned models
        fine_tuned_models = BaseSearchModel.list_fine_tuned_models()

        if fine_tuned_models:
            print("Available fine-tuned models:")
            for name, path in fine_tuned_models.items():
                print(f"- {name}: {path}")
        else:
            print("No fine-tuned models found.")
        ```
    """
    fine_tuned_models: Dict[str, str] = {}
    model_dir = "model/fine-tuned"

    if not os.path.exists(model_dir):
        logger.warning("Fine-tuned model directory not found: %s", model_dir)
        return fine_tuned_models

    for model_name in os.listdir(model_dir):
        model_path = os.path.join(model_dir, model_name)
        if os.path.isdir(model_path) and os.path.exists(
            os.path.join(model_path, "config.json")
        ):
            fine_tuned_models[model_name] = model_path

    return fine_tuned_models

@handle_exceptions(log_exceptions=True, include_exc_info=True)
def search(
    self,
    query: str,
    num_results: int = NUM_RESULTS,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> List[Dict[str, Any]]:
    """
    Search for entries matching the provided description or query.

    This method performs semantic search across the dataset by computing similarity
    scores between the user query and all synopses in the dataset. It returns the
    top matches sorted by relevance score.

    The search process includes:

    1. Creating sentence pairs between the query and all synopses
    2. Computing relevance scores using the cross-encoder model in batches
    3. Sorting results by score and returning the top matches

    Args:
        query: The search query or description to match against synopses.
            This should be a descriptive text that captures the content
            the user is looking for.
        num_results: Number of top matches to return, sorted by relevance score.
            Defaults to the value specified in constants.NUM_RESULTS.
        batch_size: Number of sentence pairs to process at once with the model.
            Using batches helps manage memory usage with large datasets.
            Defaults to the value specified in constants.DEFAULT_BATCH_SIZE.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries, each containing:
            - id: The entry ID from the id_column specified during initialization
            - title: The entry title
            - score: The relevance score (higher is better)
            - synopsis: A preview of the entry synopsis (truncated to 500 chars)

            The list is sorted by score in descending order.

    Raises:
        ValueError: If the query is empty or consists only of whitespace

    Example:
        ```python
        # Initialize a search model
        search_model = BaseSearchModel(
            dataset_path="data/merged_anime_dataset.csv",
            id_column="anime_id"
        )

        # Search for content about time travel
        results = search_model.search(
            query="A story where characters travel through time and change history",
            num_results=5,
            batch_size=64
        )

        # Process the top results
        for result in results:
            print(f"{result['title']} (Score: {result['score']:.2f})")
            print(f"Synopsis: {result['synopsis'][:100]}...")
        ```
    """
    if not query.strip():
        raise ValueError("Search query cannot be empty")

    logger.info("Searching for: %s", query)

    # Prepare pairs for cross-encoder scoring
    all_synopses = self.df["combined_synopsis"].tolist()
    sentence_pairs = [(query, text) for text in all_synopses]

    # Calculate total number of pairs and batches
    total_pairs = len(sentence_pairs)
    if batch_size <= 0:
        batch_size = DEFAULT_BATCH_SIZE
        logger.warning(
            "Invalid batch size provided, using default: %d", DEFAULT_BATCH_SIZE
        )

    # Compute relevance scores in batches
    logger.info(
        "Computing relevance scores with cross-encoder (batch size: %d)", batch_size
    )
    scores: List[Any] = []

    # Use tqdm to display progress for all datasets
    with tqdm(
        total=total_pairs,
        desc="Scoring",
        disable=False,
    ) as pbar:
        for i in range(0, total_pairs, batch_size):
            batch = sentence_pairs[i : i + batch_size]
            # Disable progress_bar in the model's predict method to avoid multiple progress bars
            batch_scores = self.model.predict(batch, show_progress_bar=False)

            if isinstance(batch_scores, np.ndarray):
                scores.extend(batch_scores.tolist())
            else:
                scores.extend(batch_scores)

            pbar.update(len(batch))

    # Convert scores to numpy array
    scores_array = np.array(scores)

    # Get indices of top scores
    top_indices = scores_array.argsort()[-num_results:][::-1]

    # Prepare results
    results = []
    for idx in top_indices:
        entry = self.df.iloc[idx]
        synopsis = entry["combined_synopsis"]
        results.append(
            {
                "id": entry[self.id_col],
                "title": entry["title"],
                "score": float(scores_array[idx]),
                "synopsis": (
                    synopsis[:500] + "..." if len(synopsis) > 500 else synopsis
                ),
            }
        )

    logger.info("Found %d matches", len(results))
    return results

def search(
    self,
    query: str,
    num_results: int = NUM_RESULTS,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> List[Dict[str, Any]]:
    """
    Search for manga entries matching the provided description or query.

    This method computes semantic similarity scores between the provided query
    and all manga synopses in the dataset (after optional filtering), returning
    the top matches sorted by relevance.

    The search process includes:

    1. Optional filtering of the dataset (e.g., removing light novels)
    2. Creating sentence pairs between the query and all manga synopses
    3. Computing relevance scores using the cross-encoder model in batches
    4. Sorting results by score and returning the top matches

    Args:
        query: The search query or description to match against manga synopses.
            This should be a descriptive text that captures the manga content
            the user is looking for.
        num_results: Number of top matches to return, sorted by relevance score.
            Defaults to the value specified in constants.NUM_RESULTS.
        batch_size: Number of sentence pairs to process at once with the model.
            Using batches helps manage memory usage with large datasets.
            Defaults to the value specified in constants.DEFAULT_BATCH_SIZE.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries, each containing:
            - id (int): The manga ID
            - title (str): The manga title
            - score (float): The relevance score (higher is better)
            - synopsis (str): A preview of the manga synopsis (truncated to 500 chars)

            The list is sorted by score in descending order.

    Raises:
        ValueError: If the query is empty or consists only of whitespace

    Example:
        ```python
        # Search for manga about time travel
        results = manga_model.search(
            query="A story about characters who can travel through time",
            num_results=3,
            batch_size=64
        )

        # Process the top results
        for result in results:
            print(f"{result['title']} (Score: {result['score']:.2f})")
        ```
    """
    if not query.strip():
        raise ValueError("Search query cannot be empty")

    logger.info("Searching for: %s", query)

    # Get the appropriate dataframe for search
    search_df = self._get_search_dataframe()

    # Prepare pairs for cross-encoder scoring
    all_synopses = search_df["combined_synopsis"].tolist()
    sentence_pairs = [(query, text) for text in all_synopses]

    # Calculate total number of pairs and batches
    total_pairs = len(sentence_pairs)
    if batch_size <= 0:
        batch_size = DEFAULT_BATCH_SIZE
        logger.warning(
            "Invalid batch size provided, using default: %d", DEFAULT_BATCH_SIZE
        )

    # Compute relevance scores in batches
    logger.info(
        "Computing relevance scores with cross-encoder (batch size: %d)", batch_size
    )
    scores: List[Any] = []

    with tqdm(
        total=total_pairs,
        desc="Scoring",
        disable=False,
    ) as pbar:
        for i in range(0, total_pairs, batch_size):
            batch = sentence_pairs[i : i + batch_size]
            # Disable progress_bar in the model's predict method to avoid multiple progress bars
            batch_scores = self.model.predict(batch, show_progress_bar=False)

            if isinstance(batch_scores, np.ndarray):
                scores.extend(batch_scores.tolist())
            else:
                scores.extend(batch_scores)

            pbar.update(len(batch))

    scores_array = np.array(scores)

    # Get indices of top scores
    top_indices = scores_array.argsort()[-num_results:][::-1]

    # Prepare results
    results = []
    for idx in top_indices:
        entry = search_df.iloc[idx]
        synopsis = entry["combined_synopsis"]
        results.append(
            {
                "id": entry[self.id_col],
                "title": entry["title"],
                "score": float(scores_array[idx]),
                "synopsis": (
                    synopsis[:500] + "..." if len(synopsis) > 500 else synopsis
                ),
            }
        )

    logger.info("Found %d matches", len(results))
    return results