Views
No views yet
pip install transformers1# pip install transformers
2from transformers import AutoModelForCausalLM, AutoTokenizer
3checkpoint = "HuggingFaceTB/SmolLM2-1.7B"
4device = "cuda" # for GPU usage or "cpu" for CPU usage
5tokenizer = AutoTokenizer.from_pretrained(checkpoint)
6# for multiple GPUs install accelerate and do `model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")`
7model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
8inputs = tokenizer.encode("Gravity is", return_tensors="pt").to(device)
9outputs = model.generate(inputs)
10print(tokenizer.decode(outputs[0]))torch.bfloat161# pip install accelerate
2# for fp16 use `torch_dtype=torch.float16` instead
3model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16)
4inputs = tokenizer.encode("Gravity is", return_tensors="pt").to("cuda")
5outputs = model.generate(inputs)
6print(tokenizer.decode(outputs[0]))1>>> print(f"Memory footprint: {model.get_memory_footprint() / 1e6:.2f} MB")
2Memory footprint: 3422.76 MB