25 Ch. 4.4: Using Prebuilt Models from Hugging Face Hub – Automatic Text Generation using the Llama 3.2 Model

export HF_TOKEN=<the-hugging-face-token> 

#!/bin/bash
set -e

# Default environment name
ENVNAME=".llama-env" 

# Default location for the environment
ENVLOC="$SLURM_TMPDIR" 
[ -n "$ENVLOC" ] || ENVLOC="$HOME"

# Parse arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --envname|-e) ENVNAME="$2"; shift ;;
        --envloc|-l) ENVLOC="$2"; shift ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

# Load the required modules
module load cuda cudnn python

# Create the virtual environment
virtualenv --no-download --seeder=pip $ENVLOC/$ENVNAME

# Activate the virtual environment
source $ENVLOC/$ENVNAME/bin/activate

# Upgrade pip
pip install --no-index --upgrade pip

# Install the required packages
pip install --no-index transformers torch

# Export the environment variables
export ENVNAME
export ENVLOC

#!/bin/bash
# drac-llama-download.sh
set -e

# Set the hugging face token
HF_TOKEN=""

# Directory to save the model
MODELDIR="$HOME/models" 

# Model name
LLAMA_MODEL="Llama-3.2-3B" 

# Parse keyword arguments for hugging face token
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -t|--hf-token) 
            if [[ -z "$2" ]]; then
            echo "Error: --hf-token requires a value."
            exit 1
            fi
            HF_TOKEN="$2"; shift ;;
        -m|--modeldir) MODELDIR="$2"; shift ;;
        -n|--modelname) LLAMA_MODEL="$2"; shift ;;
        -h|--help)
            echo "Usage: $0 --hf-token <hugging_face_token>"
            echo ""
            echo "Options:"
            echo "  -t, --hf-token    Hugging Face token for authentication (required)"
            echo "  -m, --modeldir    Path to the directory where models are stored (default: $HOME/models)"
            echo "  -n, --modelname   Name of the Llama model to download (default: Llama-3.2-3B)"
            echo "  -h, --help        Display this help message"
            exit 0
            ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

# Ensure the token are provided
if [[ -z "$HF_TOKEN" ]]; then
    echo "Error: --hf-token is required."
    exit 1
fi

# The name of the environment to create
ENVNAME="llama-env-download" 

# Create the environment and install basic dependencies
source llama-env.sh --envname $ENVNAME --envloc $HOME

echo "HF_TOKEN: $HF_TOKEN"
echo "MODELDIR: $MODELDIR"
echo "LLAMA_MODEL: $LLAMA_MODEL"
echo "ENVNAME: $ENVNAME"

# Install the huggingface-hub CLI
pip install --upgrade huggingface-hub[cli]

# Download the model using huggingface-cli
# huggingface-cli download meta-llama/$LLAMA_MODEL --include "original/*" --local-dir $MODELDIR/$LLAMA_MODEL --token $HF_TOKEN 
huggingface-cli download meta-llama/$LLAMA_MODEL --local-dir $MODELDIR/$LLAMA_MODEL --token $HF_TOKEN 

# Clean up
deactivate
rm -rf $ENVLOC/$ENVNAME
echo "Finished successfully"

export MODEL_LOCATION=$HOME/models

export MODEL_LOCATION=$SCRATCH/models

bash drac-llama-download.sh --modeldir $MODEL_LOCATION --modelname <model-name> --hf-token $HF_TOKEN

bash drac-llama-download.sh --modeldir $MODEL_LOCATION --modelname Llama-3.2-3B --hf-token $HF_TOKEN

bash drac-llama-download.sh --modeldir $MODEL_LOCATION --modelname Llama-3.2-3B-Instruct --hf-token $HF_TOKEN

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
import argparse
from transformers import pipeline
import json


homepath = os.getenv("HOME")

parser = argparse.ArgumentParser(description="Generate text using a LLaMA model.")
parser.add_argument("-m","--models_path",type=str,default=f"{homepath}/models",help="Path to the model directory.")
parser.add_argument("-n","--model_name",type=str,default="Llama-3.2-3B",help="Name of the model to use.")
parser.add_argument("-p","--prompt",type=str,default="Tell a story in more than 300 but less than 500 words. The story should start with: Once upon a time, in a digital research alliance cluster, a group of researchers embarked on a quest to harness the power of artificial intelligence. They were",help="Prompt to start the text generation.")
parser.add_argument("-s","--source_prompts_file",type=str,default=None,help="Path to the source prompt file. File should be a list with a list of dictionaries with keys 'role' and 'content'.")
parser.add_argument("-o","--output_file",type=str,default=None,help="Path to the output file to save the generated text.")
parser.add_argument("-l","--max_new_tokens",type=int,default=1024,help="Maximum number of new tokens to generate.")

def main():
    print("Drac LLaMA Text Generation")
    print("===================================")
    args = parser.parse_args()
    print(args) 

    pipe = pipeline(
        "text-generation", 
        model=os.path.join(args.models_path, args.model_name),
        torch_dtype=torch.bfloat16, 
    )
    
    output_content = []
    
    if args.source_prompts_file is not None:
        # Load the source prompt file
        with open(args.source_prompts_file, 'r') as f:
            source_prompts = json.load(f)
            
        output_content = pipe(source_prompts, max_new_tokens=args.max_new_tokens)
        print(f"Input Prompt: {source_prompts}")
        print("Generated Text:")
        print(output_content)

    else:
        output_content = pipe(args.prompt, max_new_tokens=args.max_new_tokens, num_return_sequences=1, return_full_text=False)
        print(f"Input Prompt: {args.prompt}")
        print("Generated Text:")
        print(output_content)
    
    if args.output_file is not None:
        try:
            with open(args.output_file, 'w', encoding='utf-8') as f:
                json.dump(output_content, f, ensure_ascii=False, indent=4)
                print(f"Output saved to {args.output_file}")
        except Exception as e:
            print(f"Error writing to output file: {e}")


if __name__ == "__main__":
    main()

[
  {
    "role": "system",
    "content": "You are a story teller that builds up on a starting sequence. You should limit the story size to be between 300 and 500 words."
  },
  {
    "role": "user",
    "content": " Once upon a time, in a digital research alliance cluster, a group of researchers embarked on a quest to harness the power of artificial intelligence. They were"
  }
]

python drac-llama.py --models_path $SCRATCH/models --model_name Llama-3.2-3B-Instruct --output_file "output_chat_template.json" --source_prompts_file "sample_input.json"

#!/bin/bash
#SBATCH --time=1:0:0
#SBATCH --mem-per-gpu=12G
#SBATCH --gpus-per-node=1
#SBATCH --nodes=1

# The model directory
MODELDIR="$HOME/models"
# The name of the model to be used (previously downloaded)
LLAMA_MODEL="Llama-3.2-3B"
# The script to be executed
SCRIPT="drac-llama.py"
# Additional parameters to be passed to the SCRIPT
SCRIPTKWARGS=""

# Parse optional keyword arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --modeldir) MODELDIR="$2"; shift ;;
        --llama-model) LLAMA_MODEL="$2"; shift ;;
        --script) SCRIPT="$2"; shift ;;
        --scriptkwargs) SCRIPTKWARGS="$2"; shift ;;
        --help)
            echo "Usage: $0 [--modeldir <path>] [--llama-model <model_name>]"
            echo "  --modeldir       Path to the directory where models are stored (default: $HOME/models)"
            echo "  --llama-model    Name of the Llama model to download (default: Llama-3.2-3B)"
            echo "  --script         Name of the Python script to execute (default: drac-llama.py)"
            echo "  --scriptkwargs   Additional arguments to pass to the script"
            exit 0 ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

# Check if the variables are set, otherwise set them to default values
[ -n "$ENVNAME" ] || ENVNAME=".llama-env"
[ -n "$ENVLOC" ] || ENVLOC="$SLURM_TMPDIR"
[ -n "$ENVLOC" ] || ENVLOC="$HOME"

# Create the environment and install basic dependencies
source llama-env.sh --envname $ENVNAME --envloc $ENVLOC

# Print information for debugging
printf -- '-%0.s' {1..100} | xargs echo 
echo "Virtual environment: $ENVLOC/$ENVNAME"
echo "Model directory: $MODELDIR"
echo "Llama model: $LLAMA_MODEL"
echo "Script: $SCRIPT"
printf -- '-%0.s' {1..100} | xargs echo 

# Run the python script with the specified keyword arguments
python $SCRIPT $SCRIPTKWARGS \
    --models_path $MODELDIR \
    --model_name $LLAMA_MODEL 

# Cleanup
deactivate
rm -rf $ENVLOC/$ENVNAME

sbatch -A def-<someuser> drac-llama.sh \
--modeldir $SCRATCH/models \
--llama-model Llama-3.2-3B-Instruct \
--scriptkwargs "--source_prompts_file sample_input.json --output_file output_chat.json --max_new_tokens 500"