inference.py

import os
import argparse

from transformers import AutoModelForCausalLM, AutoTokenizer

# 卡号指定
os.environ['HIP_VISIBLE_DEVICES'] = '0'

def infer_hf(model_path, input_text):
    ''' transformers 推理 gemma2'''
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="auto",
    )

    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

    outputs = model.generate(**input_ids, max_new_tokens=1024)
    print(tokenizer.decode(outputs[0]))


def parse_args():
    parser = argparse.ArgumentParser(description='gemma2 inference')
    parser.add_argument('--input_text',
                        default='Write me a poem about Machine Learning.',
                        help='')
    parser.add_argument('--model_path', default='/path/of/gemma2')
    return parser.parse_args()


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

    infer_hf(args.model_path, args.input_text)