# Running the model on a GPU using different precisions:- For Float16 -
#!pip install -U transformers
#!pip install "torch>=2.1.1" -U
#!pip install accelerate
#!huggingface-cli login
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("VijayShinde1996/vrs-gemma-2b-it")
model = AutoModelForCausalLM.from_pretrained("VijayShinde1996/vrs-gemma-2b-it", device_map="auto", torch_dtype=torch.float16)

input_text = "Write me a poem about Machine Learning."
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

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