text-generation-webui/extensions/superboogav2/notebook_handler.py

"""
This module is responsible for handling and modifying the notebook text.
"""
import re

import extensions.superboogav2.parameters as parameters
from extensions.superboogav2.utils import create_context_text
from modules.logging_colors import logger

from .data_processor import preprocess_text


def _remove_special_tokens(string):
    pattern = r'(<\|begin-user-input\|>|<\|end-user-input\|>|<\|injection-point\|>)'
    return re.sub(pattern, '', string)


def input_modifier_internal(string, collector, is_chat):
    # Sanity check.
    if is_chat:
        return string

    # Find the user input
    pattern = re.compile(r"<\|begin-user-input\|>(.*?)<\|end-user-input\|>", re.DOTALL)
    match = re.search(pattern, string)
    if match:
        # Preprocess the user prompt.
        user_input = match.group(1).strip()
        user_input = preprocess_text(user_input)

        logger.debug(f"Preprocessed User Input: {user_input}")

        # Get the most similar chunks
        results = collector.get_sorted_by_dist(user_input, n_results=parameters.get_chunk_count(), max_token_count=int(parameters.get_max_token_count()))

        # Make the injection
        string = string.replace('<|injection-point|>', create_context_text(results))

    return _remove_special_tokens(string)