Google Gemma 7b inference

gemma.py

import time
from transformers import AutoTokenizer, AutoModelForCausalLM, logging
from huggingface_hub import login

# Access token setup
ACCESS_TOKEN_READ = "token"
login(token=ACCESS_TOKEN_READ)

# Model setup
tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

# Input text
input_text = "Write me a poem about Quantum computing."

# Tokenization
input_ids = tokenizer(input_text, return_tensors="pt", max_length=50, truncation=True).to("cuda")  # Setting max_length and truncation

# Silence Transformers warnings
logging.set_verbosity_error()  # Silence warnings

# Timing execution
start_time = time.time()

# Model inference
outputs = model.generate(**input_ids, max_length=100)  # Setting max_length for generation

# Calculating execution time
execution_time = time.time() - start_time

# Decoding and printing output
print(tokenizer.decode(outputs[0]))

# Printing execution time
print(f"Execution time: {execution_time} seconds")