funcs/representation_model.py

import os
import re
import spaces
from bertopic.representation import LlamaCPP

from pydantic import BaseModel

from huggingface_hub import hf_hub_download
from gradio import Warning

from bertopic.representation import KeyBERTInspired, MaximalMarginalRelevance, BaseRepresentation
from funcs.prompts import phi3_prompt, phi3_start
from funcs.helper_functions import get_or_create_env_var, GPU_SPACE_DURATION


def clean_llm_output_text(text: str) -> str:
    """
    Clean LLM output text by removing special characters.
    Keeps only: letters, numbers, spaces, dashes, and apostrophes (for contractions).
    
    Args:
        text: The text to clean
        
    Returns:
        Cleaned text with special characters removed
    """
    if not text:
        return ""
    
    # Keep only alphanumeric characters, spaces, dashes, and apostrophes
    # This regex keeps: a-z, A-Z, 0-9, spaces, hyphens/dashes, and apostrophes
    cleaned = re.sub(r'[^a-zA-Z0-9\s\-\']', '', text)
    
    # Clean up multiple spaces and strip
    cleaned = re.sub(r'\s+', ' ', cleaned)
    cleaned = cleaned.strip()
    
    return cleaned


def patch_llama_create_chat_completion(llama_model):
    """
    Monkey-patch the create_chat_completion method on a Llama model instance
    to use raw completion instead of chat format handler.
    This avoids the "System role not supported" error for models like phi3.
    
    Args:
        llama_model: The Llama model instance to patch
        
    Returns:
        The same llama_model instance with patched create_chat_completion method
    """
    def patched_create_chat_completion(messages, **kwargs):
        """
        Override create_chat_completion to use raw completion.
        This avoids the chat format handler that requires system roles (not supported by phi3).
        BERTopic's LlamaCPP formats messages and uses the prompt template, so we reconstruct
        the full prompt from the messages.
        """
        # Reconstruct the prompt from messages
        # BERTopic's LlamaCPP passes messages in OpenAI format: [{"role": "user", "content": "..."}]
        prompt_parts = []
        for msg in messages:
            if isinstance(msg, dict):
                role = msg.get('role', 'user')
                content = msg.get('content', '')
                # Skip system messages as phi3 doesn't support them
                if role != 'system' and content:
                    prompt_parts.append(content)
            else:
                prompt_parts.append(str(msg))
        
        # Join all message contents into a single prompt
        prompt = '\n'.join(prompt_parts) if prompt_parts else ''
        
        # Use raw completion instead of chat completion
        # This avoids the chat format handler that requires system roles
        # Remove chat-specific kwargs that might cause issues, but enable streaming
        completion_kwargs = {k: v for k, v in kwargs.items() 
                           if k not in ['messages', 'chat_format', 'chat_handler']}
        
        # Enable streaming to show output in real-time
        completion_kwargs['stream'] = True
        
        # Use create_completion for raw text completion (not chat completion)
        # With stream=True, this returns a generator of CompletionChunk objects
        text_parts = []
        try:
            # Create completion with streaming enabled
            completion_stream = llama_model.create_completion(prompt, **completion_kwargs)
            
            # Iterate through the stream and collect text
            print("\nLLM Output: ", end="", flush=True)  # Print prefix without newline
            for chunk in completion_stream:
                # Extract text from each chunk
                chunk_text = ""
                
                # Handle dictionary chunks (the format returned by llama_cpp)
                if isinstance(chunk, dict):
                    # Extract from chunk['choices'][0]['text'] - this is the standard format
                    if 'choices' in chunk and len(chunk['choices']) > 0:
                        choice = chunk['choices'][0]
                        if isinstance(choice, dict):
                            chunk_text = choice.get('text', '') or choice.get('content', '')
                        elif hasattr(choice, 'text'):
                            chunk_text = choice.text
                        elif hasattr(choice, 'content'):
                            chunk_text = choice.content
                    elif 'text' in chunk:
                        chunk_text = chunk['text']
                    elif 'content' in chunk:
                        chunk_text = chunk['content']
                
                # Try different ways to extract text from the chunk (object format)
                elif hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                    choice = chunk.choices[0]
                    if hasattr(choice, 'text'):
                        chunk_text = choice.text
                    elif hasattr(choice, 'delta') and hasattr(choice.delta, 'content'):
                        # Some formats use delta.content
                        chunk_text = choice.delta.content or ""
                    elif hasattr(choice, 'content'):
                        chunk_text = choice.content
                    elif isinstance(choice, dict):
                        chunk_text = choice.get('text', '') or choice.get('delta', {}).get('content', '')
                elif hasattr(chunk, 'text'):
                    chunk_text = chunk.text
                elif isinstance(chunk, str):
                    chunk_text = chunk
                elif hasattr(chunk, '__dict__'):
                    # Check various possible attributes
                    chunk_dict = chunk.__dict__
                    if 'text' in chunk_dict:
                        chunk_text = chunk_dict['text']
                    elif 'choices' in chunk_dict:
                        choices = chunk_dict['choices']
                        if choices and len(choices) > 0:
                            if isinstance(choices[0], dict):
                                chunk_text = choices[0].get('text', '') or choices[0].get('delta', {}).get('content', '')
                            elif hasattr(choices[0], 'text'):
                                chunk_text = choices[0].text
                            elif hasattr(choices[0], 'delta'):
                                delta = choices[0].delta
                                if hasattr(delta, 'content'):
                                    chunk_text = delta.content or ""
                
                # Only add non-empty text and filter out debug messages
                if chunk_text and chunk_text.strip():
                    # Filter out llama.cpp debug messages
                    if not any(debug_keyword in chunk_text for debug_keyword in [
                        'llama_perf_context_print', 'Llama.generate', 'load time', 
                        'prompt eval time', 'eval time', 'total time', 'prefix-match hit'
                    ]):
                        text_parts.append(chunk_text)
                        print(chunk_text, end="", flush=True)  # Print without newline, flush immediately
            
            print()  # Newline after streaming is complete
            text = ''.join(text_parts)
            
            # Clean the text to remove special characters
            text = clean_llm_output_text(text)
            
            # If no text was collected, there might be an issue with chunk extraction
            if not text:
                print("Warning: No text extracted from streaming chunks. Chunk structure may be different.")
                print("Falling back to non-streaming mode.")
                raise Exception("No text in stream")
            
        except (AttributeError, TypeError, Exception) as e:
            # Fallback to non-streaming if create_completion doesn't exist or streaming fails
            print(f"\nStreaming failed, falling back to non-streaming mode: {e}")
            completion_kwargs.pop('stream', None)  # Remove stream parameter
            try:
                completion = llama_model.create_completion(prompt, **completion_kwargs)
            except AttributeError:
                completion = llama_model(prompt, **completion_kwargs)
            
            # Extract text from the completion object
            text = ""
            if hasattr(completion, 'choices') and len(completion.choices) > 0:
                # Standard Completion object format
                if hasattr(completion.choices[0], 'text'):
                    text = completion.choices[0].text
                elif hasattr(completion.choices[0], 'content'):
                    text = completion.choices[0].content
            elif hasattr(completion, 'text'):
                # Direct text attribute
                text = completion.text
            elif isinstance(completion, str):
                # Already a string
                text = completion
            elif hasattr(completion, '__dict__'):
                # Try to get text from object attributes
                if 'text' in completion.__dict__:
                    text = completion.__dict__['text']
                elif 'choices' in completion.__dict__:
                    choices = completion.__dict__['choices']
                    if choices and len(choices) > 0:
                        if isinstance(choices[0], dict):
                            text = choices[0].get('text', '')
                        elif hasattr(choices[0], 'text'):
                            text = choices[0].text
            else:
                # Last resort: convert to string (but this might not work well)
                text = str(completion)
        
        # Clean up the text - remove special characters and whitespace
        text = clean_llm_output_text(text) if text else ""
        
        # Create a chat completion response as a dictionary
        # BERTopic accesses it as: response["choices"][0]["message"]["content"]
        # Always return a dictionary to ensure it's subscriptable
        return {
            "choices": [{
                "message": {
                    "content": text,
                    "role": "assistant"
                },
                "finish_reason": "stop",
                "index": 0
            }],
            "id": "custom",
            "created": 0,
            "model": "",
            "object": "chat.completion"
        }
    
    # Replace the method on the instance
    llama_model.create_chat_completion = patched_create_chat_completion
    
    return llama_model

chosen_prompt = phi3_prompt #open_hermes_prompt # stablelm_prompt 
chosen_start_tag =  phi3_start #open_hermes_start # stablelm_start

random_seed = 42

RUNNING_ON_AWS = get_or_create_env_var('RUNNING_ON_AWS', '0')
print(f'The value of RUNNING_ON_AWS is {RUNNING_ON_AWS}')

USE_GPU = get_or_create_env_var('USE_GPU', '0')
print(f'The value of USE_GPU is {USE_GPU}')

# from torch import cuda, backends, version, get_num_threads

# print("Is CUDA enabled? ", cuda.is_available())
# print("Is a CUDA device available on this computer?", backends.cudnn.enabled)
# if cuda.is_available():
#     torch_device = "gpu"
#     print("Cuda version installed is: ", version.cuda)
#     high_quality_mode = "Yes"
#     os.system("nvidia-smi")
# else: 
#     torch_device =  "cpu"
#     high_quality_mode = "No"

if USE_GPU == "1":
    print("Using GPU for representation functions")
    torch_device = "gpu"
    #print("Cuda version installed is: ", version.cuda)
    high_quality_mode = "Yes"
    os.system("nvidia-smi")
else:
    print("Using CPU for representation functions")
    torch_device =  "cpu"
    high_quality_mode = "No"

# Currently set n_gpu_layers to 0 even with cuda due to persistent bugs in implementation with cuda
print("torch device for representation functions:", torch_device)
if torch_device == "gpu":
    low_resource_mode = "No"
    n_gpu_layers = -1 # i.e. all
else: #     torch_device =  "cpu"
    low_resource_mode = "Yes"
    n_gpu_layers = 0

#print("Running on device:", torch_device)
from torch import get_num_threads
n_threads = get_num_threads()
print("CPU n_threads:", n_threads)

# Default Model parameters
temperature: float = 0.1
top_k: int = 3
top_p: float = 1
repeat_penalty: float = 1.1
last_n_tokens_size: int = 128
max_tokens: int = 500
seed: int = random_seed
reset: bool = True
stream: bool = False
n_threads: int = n_threads
n_batch:int = 512
n_ctx:int = 8192 #4096. # Set to 8192 just to avoid any exceeded context window issues
sample:bool = True
trust_remote_code:bool =True

class LLamacppInitConfigGpu(BaseModel):
    last_n_tokens_size: int
    seed: int
    n_threads: int
    n_batch: int
    n_ctx: int
    n_gpu_layers: int
    temperature: float
    top_k: int
    top_p: float
    repeat_penalty: float
    max_tokens: int
    reset: bool
    stream: bool
    stop: str
    trust_remote_code:bool

    def update_gpu(self, new_value: int):
        self.n_gpu_layers = new_value

llm_config = LLamacppInitConfigGpu(last_n_tokens_size=last_n_tokens_size,
    seed=seed,
    n_threads=n_threads,
    n_batch=n_batch,
    n_ctx=n_ctx,
    n_gpu_layers=n_gpu_layers,
    temperature=temperature,
    top_k=top_k,
    top_p=top_p,
    repeat_penalty=repeat_penalty,
    max_tokens=max_tokens,
    reset=reset,
    stream=stream,
    stop=chosen_start_tag,
    trust_remote_code=trust_remote_code)

## Create representation model parameters ##
keybert = KeyBERTInspired(random_state=random_seed)
mmr = MaximalMarginalRelevance(diversity=0.5)
base_rep = BaseRepresentation()

# Find model file
def find_model_file(hf_model_name: str, hf_model_file: str, search_folder: str, sub_folder: str) -> str:
    """
    Finds the specified model file within the given search folder and subfolder.

    Args:
        hf_model_name (str): The name of the Hugging Face model.
        hf_model_file (str): The specific file name of the model to find.
        search_folder (str): The base folder to start the search.
        sub_folder (str): The subfolder within the search folder to look into.

    Returns:
        str: The path to the found model file, or None if the file is not found.
    """

    hf_loc = search_folder #os.environ["HF_HOME"]
    hf_sub_loc = search_folder + sub_folder #os.environ["HF_HOME"] 

    if sub_folder == "/hub/":
        hf_model_name_path = hf_sub_loc + 'models--' + hf_model_name.replace("/","--")
    else:
        hf_model_name_path = hf_sub_loc

    def find_file(root_folder, file_name):
        for root, dirs, files in os.walk(root_folder):
            if file_name in files:
                return os.path.join(root, file_name)
        return None

    # Example usage
    folder_path = hf_model_name_path  # Replace with your folder path
    file_to_find = hf_model_file         # Replace with the file name you're looking for

    print("Searching for model file", hf_model_file, "in:", hf_model_name_path)

    found_file = find_file(folder_path, file_to_find) # os.environ["HF_HOME"]
    
    return found_file

@spaces.GPU(duration=GPU_SPACE_DURATION)
def create_representation_model(representation_type: str, llm_config: dict, hf_model_name: str, hf_model_file: str, chosen_start_tag: str, low_resource_mode: bool) -> dict:
    """
    Creates a representation model based on the specified type and configuration.

    Args:
        representation_type (str): The type of representation model to create (e.g., "LLM", "KeyBERT").
        llm_config (dict): Configuration settings for the LLM model.
        hf_model_name (str): The name of the Hugging Face model.
        hf_model_file (str): The specific file name of the model to find.
        chosen_start_tag (str): The start tag to use for the model.
        low_resource_mode (bool): Whether to enable low resource mode.

    Returns:
        dict: A dictionary containing the created representation model.
    """

    if representation_type == "LLM":
        print("RUNNING_ON_AWS:", RUNNING_ON_AWS)
        if RUNNING_ON_AWS=="1":
            error_message = "LLM representation not available on AWS due to model size restrictions. Returning base representation"
            Warning(error_message, duration=5)
            print(error_message)
            representation_model = {"LLM":base_rep}
            return representation_model
        # Else import Llama
        else:
            from llama_cpp import Llama

        print("Generating LLM representation")
        # Use llama.cpp to load in model

        # Check for HF_HOME environment variable and supply a default value if it's not found (typical location for huggingface models)
        base_folder = "model" #"~/.cache/huggingface/hub"
        hf_home_value = os.getenv("HF_HOME", base_folder)

        # Expand the user symbol '~' to the full home directory path
        if "~" in base_folder:
            hf_home_value = os.path.expanduser(hf_home_value)

        # Check if the directory exists, create it if it doesn't
        if not os.path.exists(hf_home_value):
            os.makedirs(hf_home_value)

        print("Searching base folder for model:", hf_home_value)

        found_file = find_model_file(hf_model_name, hf_model_file,  hf_home_value, "/rep/")

        if found_file:
            print(f"Model file found in model folder: {found_file}")

        else:
            found_file = find_model_file(hf_model_name, hf_model_file,  hf_home_value, "/hub/")

        if not found_file:
            error = "File not found in HF hub directory or in local model file."
            print(error, " Downloading model from hub")

            found_file = hf_hub_download(repo_id=hf_model_name, filename=hf_model_file)#, local_dir=hf_home_value) # cache_dir

            print("Downloaded model from Huggingface Hub to: ", found_file)

        print("Loading representation model with", llm_config.n_gpu_layers, "layers allocated to GPU.")

        #llm_config.n_gpu_layers
        # Initialize Llama model - try to disable chat format handler if supported
        # This helps avoid "System role not supported" error for models like phi3
        try:
            llm = Llama(model_path=found_file, stop=chosen_start_tag, n_gpu_layers=llm_config.n_gpu_layers, n_ctx=llm_config.n_ctx, seed=seed, chat_format=None)
        except TypeError:
            # If chat_format parameter doesn't exist, try without it or with chat_handler
            try:
                llm = Llama(model_path=found_file, stop=chosen_start_tag, n_gpu_layers=llm_config.n_gpu_layers, n_ctx=llm_config.n_ctx, seed=seed, chat_handler=None)
            except TypeError:
                # Fall back to basic initialization if chat format parameters don't exist
                llm = Llama(model_path=found_file, stop=chosen_start_tag, n_gpu_layers=llm_config.n_gpu_layers, n_ctx=llm_config.n_ctx, seed=seed)
        
        # Monkey-patch the create_chat_completion method to use raw completion
        # This avoids the chat format handler that requires system roles (not supported by phi3)
        # We patch the instance directly so it still passes isinstance checks in BERTopic
        llm = patch_llama_create_chat_completion(llm)
        
        #print(llm.n_gpu_layers)
        #print("Chosen prompt:", chosen_prompt)
        llm_model = LlamaCPP(llm, prompt=chosen_prompt)#, **gen_config.model_dump())

        # All representation models
        representation_model = {
        "LLM": llm_model
        }

    elif representation_type == "KeyBERT":
        print("Generating KeyBERT representation")
        #representation_model = {"mmr": mmr}
        representation_model = {"KeyBERT": keybert}

    elif representation_type == "MMR":
        print("Generating MMR representation")
        representation_model = {"MMR": mmr}

    else:
        print("Generating default representation type")
        representation_model = {"Default":base_rep}
        
    return representation_model