Upload README.md with huggingface_hub

README.md

@@ -66,16 +66,43 @@ No custom Gemma 3 modeling code is used.
 import torch
 from transformers import Gemma3ForCausalLM, PreTrainedTokenizerFast
 
-model_dir = "hf"
-tokenizer = PreTrainedTokenizerFast.from_pretrained(model_dir)
-model = Gemma3ForCausalLM.from_pretrained(model_dir)
-model.eval()
-
-ids = [tokenizer.bos_token_id] + tokenizer.encode("Once upon", add_special_tokens=False)
-input_ids = torch.tensor([ids], dtype=torch.long)
-
-with torch.no_grad():
-    out = model.generate(input_ids=input_ids, max_new_tokens=64, do_sample=False)
-
-print(tokenizer.decode(out[0], skip_special_tokens=True))
+def main():
+    repo_id = "shibatch/tinygemma3-2m"
+
+    print("Loading tokenizer...")
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(repo_id, subfolder="hf")
+
+    print("Loading Gemma3 model weights...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model = Gemma3ForCausalLM.from_pretrained(
+        repo_id,
+        subfolder="hf",
+        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+    ).to(device)
+    model.eval()
+
+    prompt = "Once upon"
+    print(f"\nInput prompt: {prompt}")
+
+    input_ids = tokenizer.encode(prompt, add_special_tokens=False)
+    input_ids = [tokenizer.bos_token_id] + input_ids
+    input_ids = torch.tensor([input_ids], dtype=torch.long, device=device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids,
+            max_new_tokens=100,
+            do_sample=False,
+            repetition_penalty=1.0,
+            top_p=1.0,
+            pad_token_id=tokenizer.pad_token_id or tokenizer.bos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    print(f"Generated output: {generated_text}")
+
+if __name__ == "__main__":
+    main()
 ```