from transformers import BitsAndBytesConfig


import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os


# Function to convert a string representation of truth to true (1) or false (0)
def boolean_string(s):

    """
    Convert a string representation of truth to true (1) or false (0).
    Args:
        s (str): The string to convert, should be "true" or "false" (case insensitive).
    Returns:
        bool: True if the string is "true", False if it is "false".
    """

    if s.lower() not in {"false", "true"}:
        raise ValueError("Not a valid boolean string")
    
    return s.lower() == "true"

# Function to load a pre-trained model and tokenizer with low CPU memory usage and 8-bit quantization
def load_model(path):

    """
    Load a pre-trained model and tokenizer from the specified path with low CPU memory usage and 8-bit quantization.
    Args:
        path (str): The path to the pre-trained model.
    Returns:
        model: The loaded model.
        tokenizer: The loaded tokenizer.
    """

    # Get the hugging face token from the environment variable
    auth_token = os.environ["HUGGINGFACE_TOKEN"]

    # Create a BitsAndBytesConfig for 8-bit precision
    bnb_config = BitsAndBytesConfig(
        load_in_8bit=True,  # Enable 8-bit quantization
        compute_dtype=torch.float16  # Optionally set the compute dtype to float16 for efficiency
    )

    print("Loading model...")
    
    # Determining the tokenizer based on the model path
    if "llama" in path:
        tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, token=auth_token)
    elif "mistral" in path:
        tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True, token=auth_token)
    else:
        tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, token=auth_token)

    print("Tokenizer loaded.")
    
    # Load the model with low CPU memory usage and 8-bit quantization
    model = AutoModelForCausalLM.from_pretrained(
        path, low_cpu_mem_usage=True, torch_dtype=torch.float16, quantization_config=bnb_config, token=auth_token
    )
    
    print("Model loaded.")

    return model, tokenizer
