Views
No views yet
1import torch
2from transformers import AutoTokenizer
3from huggingface_hub import hf_hub_download
4from SmolLm3 import LlamaModel
5import yaml
6# Download the model file
7model_path = hf_hub_download(
8 repo_id="crpatel/SmolLM2-135M-cosmopedia2-70kSteps",
9 filename="model.pt"
10)
11
12config = yaml.load(open('config_smollm2_135M.yaml', "r"), Loader=yaml.FullLoader)
13model = LlamaModel(config['model'])
14model.load_state_dict(torch.load(model_path, map_location='cpu'))
15tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
16# cpu = torch.device('cpu')
17encoded_text = tokenizer.encode('Once Upon time ', return_tensors="pt").to('cpu')
18print(encoded_text)
19generated_text2=model.generate(idx=encoded_text, max_new_tokens=100, context_length=50,
20 temperature=0.9,
21 top_k=2, eos_token=tokenizer.eos_token_id,
22 device='cpu')
23decoded_text2=tokenizer.decode(generated_text2.squeeze(0))
24print(decoded_text2)