Move to in-library checkpoint

README.md

@@ -27,9 +27,9 @@ For full details of this model please read the [white paper](https://arxiv.org/a
 
 ## Usage
 ### Presequities
-Jamba requires you use `transformers` version 4.39.0 or higher:
+In order to use Jamba, it is recommended you use `transformers` version 4.40.0 or higher (usage of version 4.39.0 is required):
 ```bash
-pip install transformers>=4.39.0
+pip install transformers>=4.40.0
 ```
 
 In order to run optimized Mamba implementations, you first need to install `mamba-ssm` and `causal-conv1d`:
@@ -41,12 +41,10 @@ You also have to have the model on a CUDA device.
 You can run the model not using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly lower latencies. In order to do that, you'll need to specify `use_mamba_kernels=False` when loading the model.
 
 ### Run the model
-Please note that, at the moment, `trust_remote_code=True` is required for running the new Jamba architecture.
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
-                                             trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
 tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
 
 input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]
@@ -57,6 +55,8 @@ print(tokenizer.batch_decode(outputs))
 # ["<|startoftext|>In the recent Super Bowl LVIII, the Kansas City Chiefs emerged victorious, defeating the San Francisco 49ers in a thrilling overtime showdown. The game was a nail-biter, with both teams showcasing their skills and determination.\n\nThe Chiefs, led by their star quarterback Patrick Mahomes, displayed their offensive prowess, while the 49ers, led by their strong defense, put up a tough fight. The game went into overtime, with the Chiefs ultimately securing the win with a touchdown.\n\nThe victory marked the Chiefs' second Super Bowl win in four years, solidifying their status as one of the top teams in the NFL. The game was a testament to the skill and talent of both teams, and a thrilling end to the NFL season.\n\nThe Super Bowl is not just about the game itself, but also about the halftime show and the commercials. This year's halftime show featured a star-studded lineup, including Usher, Alicia Keys, and Lil Jon. The show was a spectacle of music and dance, with the performers delivering an energetic and entertaining performance.\n"]
 ```
 
+Please note that if you're using `transformers<4.40.0`, `trust_remote_code=True` is required for running the new Jamba architecture.
+
 <details>
 <summary><strong>Loading the model in half precision</strong></summary>
   
@@ -66,7 +66,6 @@ print(tokenizer.batch_decode(outputs))
 from transformers import AutoModelForCausalLM
 import torch
 model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
-                                             trust_remote_code=True,
                                              torch_dtype=torch.bfloat16)    # you can also use torch_dtype=torch.float16
 ```
 
@@ -75,7 +74,6 @@ When using half precision, you can enable the [FlashAttention2](https://github.c
 from transformers import AutoModelForCausalLM
 import torch
 model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
-                                             trust_remote_code=True,
                                              torch_dtype=torch.bfloat16,
                                              attn_implementation="flash_attention_2",
                                              device_map="auto")
@@ -91,7 +89,6 @@ from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 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)
@@ -108,7 +105,7 @@ from peft import LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 
 tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", trust_remote_code=True, device_map='auto')
+model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map='auto')
 
 dataset = load_dataset("Abirate/english_quotes", split="train")
 training_args = TrainingArguments(