Update README.md

README.md

@@ -11,7 +11,7 @@ tags:
 - trl
 ---
 
-# Athena-gemma-2-2b-lt for coding
+# Athena-codegemma-2-2b-lt for coding
 
 Supervised fine tuned (sft unsloth) for coding with EpistemeAI coding dataset.
 
@@ -49,7 +49,7 @@ import torch
 from transformers import pipeline
 pipe = pipeline(
     "text-generation",
-    model="EpistemeAI/Athena-gemma-2-2b-it",
+    model="EpistemeAI/Athena-codegemma-2-2b-it",
     model_kwargs={"torch_dtype": torch.bfloat16},
     device="cuda",  # replace with "mps" to run on a Mac device
 )
@@ -70,7 +70,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
 model = AutoModelForCausalLM.from_pretrained(
-    "EpistemeAI/Athena-gemma-2-2b-it",
+    "EpistemeAI/Athena-codegemma-2-2b-it",
     device_map="auto",
     torch_dtype=torch.bfloat16,
 )
@@ -104,7 +104,7 @@ You can also use `float32` if you skip the dtype, but no precision increase will
 from transformers import AutoTokenizer, AutoModelForCausalLM
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
 model = AutoModelForCausalLM.from_pretrained(
-    "EpistemeAI/Athena-gemma-2-2b-it",
+    "EpistemeAI/Athena-codegemma-2-2b-it",
     device_map="auto",
 )
 input_text = "Write me a poem about Machine Learning."
@@ -135,7 +135,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
 model = AutoModelForCausalLM.from_pretrained(
-    "EpistemeAI/Athena-gemma-2-2b-it",
+    "EpistemeAI/Athena-codegemma-2-2b-it",
     quantization_config=quantization_config,
 )
 input_text = "Write me a poem about Machine Learning."
@@ -155,7 +155,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 quantization_config = BitsAndBytesConfig(load_in_4bit=True)
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
 model = AutoModelForCausalLM.from_pretrained(
-    "EpistemeAI/Athena-gemma-2-2b-it",
+    "EpistemeAI/Athena-codegemma-2-2b-it",
     quantization_config=quantization_config,
 )
 input_text = "Write me a poem about Machine Learning."
@@ -185,7 +185,7 @@ import torch
 torch.set_float32_matmul_precision("high")
 # load the model + tokenizer
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
-model = Gemma2ForCausalLM.from_pretrained("EpistemeAI/Athena-gemma-2-2b-it", torch_dtype=torch.bfloat16)
+model = Gemma2ForCausalLM.from_pretrained("EpistemeAI/Athena-codegemma-2-2b-it", torch_dtype=torch.bfloat16)
 model.to("cuda")
 # apply the torch compile transformation
 model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)