1
0

trying and failiing to get olmo to run locally

This commit is contained in:
Matthew Gaughan 2025-05-05 10:59:02 -07:00
parent ed78364422
commit 1bcf2d8e72

25
test.py Normal file
View File

@ -0,0 +1,25 @@
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
olmo = AutoModelForCausalLM.from_pretrained(
"allenai/OLMo-2-1124-7B-hf",
torch_dtype=torch.float32,
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-1124-7B-hf")
message = ["Honolulu is a "]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
response = olmo.generate(
**inputs,
max_new_tokens=128,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.5
)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])