Create utils.py

utils.py

@@ -0,0 +1,181 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+import re
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer
+from peft import PeftModel
+
+def get_device_map():
+    num_gpus = torch.cuda.device_count()
+
+    if num_gpus > 1:
+        print("More than one GPU found. Setting device_map to use CUDA device 0.")
+        return 'cuda:0'
+    else:
+        return 'auto'
+
+def check_adapter_path(adapters_name):
+    """
+    Checks if the adapter path is correctly set and not a placeholder.
+    Args:
+    adapters_name (str): The file path for the adapters.
+    Raises:
+    ValueError: If the adapters_name contains placeholder characters.
+    """
+    if '<' in adapters_name or '>' in adapters_name:
+        raise ValueError("The adapter path has not been set correctly.")
+
+def load_tokenizer(model_name):
+    """
+    Loads and returns a tokenizer for the specified model.
+    Args:
+    model_name (str): The name of the model for which to load the tokenizer.
+    Returns:
+    AutoTokenizer: The loaded tokenizer with special tokens added and padding side set.
+    """
+    tok = AutoTokenizer.from_pretrained(model_name, device_map=get_device_map(), trust_remote_code=True)
+    tok.add_special_tokens({'pad_token': '[PAD]'})
+    tok.padding_side = 'right'  # TRL requires right padding
+    return tok
+
+def load_model(model_name, torch_dtype, quant_type):
+    """
+    Loads and returns a model with the specified quantization configuration.
+    If more than one GPU is available, wraps the model with DataParallel.
+    Args:
+    model_name (str): The name of the model to load.
+    torch_dtype (torch.dtype): The data type for model weights (e.g., torch.float16).
+    quant_type (str): The quantization type to use.
+    Returns:
+    AutoModelForCausalLM: The loaded model possibly wrapped with DataParallel.
+    """
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path=model_name,
+            trust_remote_code=True,
+            device_map=get_device_map(),
+            torch_dtype=torch_dtype,
+            quantization_config=BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch_dtype,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type=quant_type
+            ),
+        )
+      
+        return model
+    except Exception as e:
+        raise RuntimeError(f"Error loading model: {e}")
+
+def resize_embeddings(model, tokenizer):
+    """
+    Resizes the token embeddings in the model to account for new tokens.
+    Args:
+    model (AutoModelForCausalLM): The model whose token embeddings will be resized.
+    tokenizer (AutoTokenizer): The tokenizer corresponding to the model.
+    """
+    model.resize_token_embeddings(len(tokenizer))
+
+def load_peft_model(model, adapters_name):
+    """
+    Loads the PEFT model from the pretrained model and specified adapters.
+    Args:
+    model (AutoModelForCausalLM): The base model.
+    adapters_name (str): Path to the adapters file.
+    Returns:
+    PeftModel: The PEFT model with the loaded adapters.
+    """
+    return PeftModel.from_pretrained(model, adapters_name)
+
+def get_device():
+    """
+    Determines and returns the device to use for computations.
+    If CUDA is available, returns a CUDA device, otherwise returns a CPU device.
+    Prints the number of GPUs available if CUDA is used.
+    Returns:
+    torch.device: The device to use.
+    """
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        print(f"Number of GPUs available: {torch.cuda.device_count()}")
+    else:
+        device = torch.device("cpu")
+    return device
+
+def run_prompt(model, tokenizer, device, template):
+    """
+    Runs an interactive prompt where the user can enter text to get generated responses.
+    Continues to prompt the user for input until '#end' is entered.
+    Args:
+    model (AutoModelForCausalLM): The model to use for text generation.
+    tokenizer (AutoTokenizer): The tokenizer to use for encoding the input text.
+    device (torch.device): The device on which to perform the computation.
+    template (str): The template string to format the input text.
+    """
+    while True:
+        new_input = input("Enter your text (type #end to stop): ")
+        if new_input == "#end":
+            break
+
+        try:
+            _ = generate_text(model, tokenizer, device, new_input, template)
+        except Exception as e:
+            print(f"An error occurred during text generation: {e}")
+            
+def generate_text(model, tokenizer, device, input_text, template):
+    """
+    Generates and returns text using the provided model and tokenizer for the input text.
+    Args:
+    model (AutoModelForCausalLM): The model to use for text generation.
+    tokenizer (AutoTokenizer): The tokenizer to use for encoding the input text.
+    device (torch.device): The device on which to perform the computation.
+    input_text (str): The input text to generate responses for.
+    template (str): The template string to format the input text.
+    Returns:
+    torch.Tensor: The generated text tensor.
+    """
+    inputs = tokenizer(template.format(input_text), return_tensors="pt")
+    inputs = inputs.to(device)  # Move input tensors to the device
+    streamer = TextStreamer(tokenizer)
+    return model.generate(**inputs, streamer=streamer,
+                          max_new_tokens=1024,
+                          pad_token_id=tokenizer.pad_token_id,
+                          eos_token_id=tokenizer.eos_token_id)
+
+def get_last_folder_alphabetically(directory_path):
+    """
+    Finds the last folder alphabetically in a specified directory.
+
+    Args:
+        directory_path (str): The path to the directory.
+
+    Returns:
+        str: The path to the last folder found alphabetically.
+        If the directory does not exist or contains no folders, a descriptive string is returned.
+    """
+    if not os.path.exists(directory_path):
+        return "Directory does not exist."
+
+    all_files_and_folders = os.listdir(directory_path)
+    only_folders = [f for f in all_files_and_folders if os.path.isdir(os.path.join(directory_path, f))]
+    if not only_folders:
+        return "No folders found in the directory."
+
+    only_folders.sort(key=natural_sort_key)
+    last_folder = only_folders[-1]
+    return os.path.join(directory_path, last_folder)
+
+def natural_sort_key(s):
+    """
+    Generates a key for sorting strings that contain numbers where the numbers should be sorted numerically,
+    and the rest alphabetically.
+
+    Args:
+        s (str): The string to be sorted.
+
+    Returns:
+        list: A list of strings and integers derived from the input string.
+    """
+    return [int(text) if text.isdigit() else text.lower() for text in re.split('([0-9]+)', s)]