pipe.py

import argparse

from huggingface_hub.constants import HF_HUB_CACHE
from transformers import pipeline, AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from dotenv import load_dotenv
load_dotenv(override=True)


def generate(
    model_id: str,
    prompt: str,
    max_tokens: int,
    load_in_4bit: bool,
    load_in_8bit: bool,
    trust_remote_code: bool
):
    if load_in_4bit and load_in_8bit:
        raise ValueError("Cannot load in both 4-bit and 8-bit")

    q_config: BitsAndBytesConfig = None

    if load_in_4bit or load_in_8bit:
        q_config = BitsAndBytesConfig(
            load_in_4bit=load_in_4bit,
            load_in_8bit=load_in_8bit
        )

    config = AutoConfig.from_pretrained(model_id, cache_dir=HF_HUB_CACHE)
    print(config)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=q_config,
        low_cpu_mem_usage=True,
        trust_remote_code=trust_remote_code,
        cache_dir=HF_HUB_CACHE
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=HF_HUB_CACHE)

    pipe = pipeline('text-generation', model=model, tokenizer=tokenizer)
    output = pipe(prompt, max_length=max_tokens, return_full_text=False)
    print(output[0]["generated_text"])


def parse_args():
    parser = argparse.ArgumentParser(
        description="Using PyTorch and transformers to generate text"
    )
    parser.add_argument(
        "-m", "--model",
        type=str,
        required=True,
        help="The model ID to use for generation"
    )
    parser.add_argument(
        "-p", "--prompt",
        type=str,
        required=True,
        help="The prompt to use for generation"
    )
    parser.add_argument(
        "-t", "--max_tokens",
        type=int,
        required=True,
        help="The maximum number of tokens to generate"
    )
    parser.add_argument(
        "--load_in_4bit",
        action="store_true",
        help="Load the model in 4-bit"
    )
    parser.add_argument(
        "--load_in_8bit",
        action="store_true",
        help="Load the model in 8-bit"
    )
    parser.add_argument(
        "--trust_remote_code",
        action="store_true",
        help="Trust remote code from Hugging Face"
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    model_id: str = args.model
    prompt: str = args.prompt
    max_tokens: int = args.max_tokens
    load_in_4bit: bool = args.load_in_4bit
    load_in_8bit: bool = args.load_in_8bit
    trust_remote_code: bool = args.trust_remote_code

    generate(
        model_id,
        prompt,
        max_tokens,
        load_in_4bit,
        load_in_8bit,
        trust_remote_code
    )