You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
38 lines
1.3 KiB
38 lines
1.3 KiB
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
import time
|
|
|
|
t_start = time.time()
|
|
|
|
# model_name = "NousResearch/Llama-2-7b-hf" # will cache on C:\Users\ftobler\.cache\huggingface\hub
|
|
model_name = "NousResearch/Hermes-3-Llama-3.2-3B" # will cache on C:\Users\ftobler\.cache\huggingface\hub
|
|
# "meta-llama/Llama-2-7b-hf" # Replace with your chosen model
|
|
|
|
# Load the model with quantization (optional)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_name,
|
|
# device_map="auto", # Automatically places parts of the model on GPU/CPU
|
|
device_map="cpu", # Automatically places parts of the model on GPU/CPU
|
|
load_in_8bit=False, # Enables 8-bit quantization if bitsandbytes is installed
|
|
)
|
|
|
|
# Load tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
|
|
print("load took %.3fs" % (time.time() - t_start))
|
|
t_start = time.time()
|
|
|
|
# Generate text
|
|
input_text = "Hello, who are you?"
|
|
inputs = tokenizer(input_text, return_tensors="pt").to("cpu") # .to("cuda") .to("cpu")
|
|
outputs = model.generate(
|
|
inputs["input_ids"],
|
|
# max_length=200,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
eos_token_id=tokenizer.eos_token_id
|
|
)
|
|
|
|
# Decode and print result
|
|
print(tokenizer.decode(outputs[0], skip_special_tokens=False))
|
|
|
|
print("genaration took %.3fs" % (time.time() - t_start))
|
|
t_start = time.time()
|