Update README.md

README.md

@@ -33,6 +33,43 @@ Since this is a base model the IKM dataset greatly affects the output. The IKM d
 ```
 ---
 
+## Inference
+
+```
+!pip install -qqq transformers>=4.39.0 mamba-ssm causal-conv1d>=1.2.0 accelerate bitsandbytes --progress-bar off
+!pip install flash-attn --no-build-isolation
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+# Load model in 8-bit precision
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_skip_modules=["mamba"]
+)
+model = AutoModelForCausalLM.from_pretrained(
+    "ai21labs/Jamba-v0.1",
+    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")
+
+# Tokenize input
+prompt = """How could we use cheese to reignite the sun? Answer:"""
+input_ids = tokenizer(
+    prompt,
+    return_tensors='pt'
+).to(model.device)["input_ids"]
+
+# Generate answer
+outputs = model.generate(input_ids, max_new_tokens=216)
+
+# Print output
+print(tokenizer.batch_decode(outputs))
+```
+
 ```
  [3731/5850 3:38:52 < 2:04:22, 0.28 it/s, Epoch 6.37/10]
 Step	Training Loss