Spaces:
Paused
Paused
File size: 2,677 Bytes
2f86d9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """Interactive Multi-turn Chat with Gemma 3.
Run this script using the configured environment:
e:\\v\\Scripts\\python.exe run_chat.py
"""
import os
import sys
# Suppress verbose TensorFlow / JAX initialization notices
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
from gemma import gm
from gemma.gm.text._tokenizer import Gemma3Tokenizer
import kagglehub
def get_weights_and_tokenizer():
"""Retrieves locally cached weights and tokenizer via KaggleHub."""
print("Checking for Gemma 3 (1B) weights...")
kaggle_id = "google/gemma-3/flax/gemma3-1b-it"
weights_dir = kagglehub.model_download(kaggle_id)
print(f"Weights directory: {weights_dir}")
# Check for inner checkpoint directory
inner_ckpt = os.path.join(weights_dir, "gemma3-1b-it")
ckpt_path = inner_ckpt if os.path.exists(inner_ckpt) else weights_dir
# Check for local tokenizer.model
tok_file = os.path.join(weights_dir, "tokenizer.model")
if os.path.exists(tok_file):
tokenizer = Gemma3Tokenizer(path=tok_file)
print("Loaded local tokenizer from disk.")
else:
tokenizer = None
return ckpt_path, tokenizer
def main():
print("=" * 60)
print("Initializing Gemma 3 (Option A: Multi-turn Chat)")
print("=" * 60)
# 1. Initialize 1B architecture
print("\nInitializing model architecture: Gemma 3 (1B)...")
model = gm.nn.Gemma3_1B()
# 2. Get local weights & tokenizer
ckpt_path, tokenizer = get_weights_and_tokenizer()
# 3. Load parameters into memory
print("\nLoading weights into memory...")
params = gm.ckpts.load_params(ckpt_path)
print("Weights loaded successfully!")
# 4. Create multi-turn chat sampler
print("\nInitializing ChatSampler (First turn will JIT-compile)...")
sampler = gm.text.ChatSampler(
model=model,
params=params,
tokenizer=tokenizer,
multi_turn=True,
print_stream=True, # Stream tokens live to console
)
print("\n" + "=" * 60)
print("Gemma is ready! Type your message (or 'exit' to quit).")
print("=" * 60 + "\n")
while True:
try:
user_input = input("\nYou: ").strip()
if not user_input:
continue
if user_input.lower() in ("exit", "quit", "q"):
print("Goodbye!")
break
print("\nGemma: ", end="", flush=True)
sampler.chat(user_input)
print()
except KeyboardInterrupt:
print("\nSession ended. Exiting.")
break
except Exception as e:
print(f"\nError during generation: {e}")
if __name__ == "__main__":
main()
|