Update README.md

README.md

@@ -42,19 +42,19 @@ Since this is a base model the IKM dataset greatly affects the output. The IKM d
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-# Load model in 8-bit precision
+# Load model in 4-bit precision
 quantization_config = BitsAndBytesConfig(
-    load_in_8bit=True,
+    load_in_4bit=True,
     llm_int8_skip_modules=["mamba"]
 )
 model = AutoModelForCausalLM.from_pretrained(
-    "ai21labs/Jamba-v0.1",
+    "Severian/Jamba-Nexus-IKM-v1",
     trust_remote_code=True,
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2",
     quantization_config=quantization_config
 )
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
+tokenizer = AutoTokenizer.from_pretrained("Severian/Jamba-Nexus-IKM-v1")
 
 # Tokenize input
 prompt = """How could we use cheese to reignite the sun? Answer:"""