Repository: huggingface/transformers
Pull request #30838 (2024-05-15 16:53:01)

Add streaming, various fixes
# What does this PR do?

- [Fixes tool pushing/loading from Hub](https://github.com/huggingface/transformers/pull/30838/commits/fed6c74c2c93a2b97e2da0fd8e4fa8a9ee3eeb84#diff-0e9df5d3d8bd1fcd6927c628729420a0ae593e2bf5c4162af579dc138aea67d1).
- [Implements streaming behaviour for `run` method](https://github.com/huggingface/transformers/pull/30838/commits/ca47a4a36ddcd824bde1b720837a569ea995e2f7)
- [Improve default prompts](https://github.com/huggingface/transformers/pull/30838/commits/050ad9473161af3d54a3d54ef1a2ffc4913070e0)

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30838). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
aymeric-roucher: @LysandreJik this is ready for review. It adds streaming to `run` method, as needed by @freddyaboulton.
freddyaboulton: The streaming api works great - thanks @aymeric-roucher !

Patch:
diff --git a/.github/workflows/model_jobs.yml b/.github/workflows/model_jobs.yml
index f88af8e39a..840df8b697 100644
--- a/.github/workflows/model_jobs.yml
+++ b/.github/workflows/model_jobs.yml
@@ -80,7 +80,7 @@ jobs:
 
       - name: Run all tests on GPU
         working-directory: /transformers
-        run: python3 -m pytest -rs -v --make-reports=${{ inputs.machine_type }}_run_models_gpu_${{ matrix.folders }}_test_reports tests/${{ matrix.folders }}
+        run: python3 -m pytest -rsfE -v --make-reports=${{ inputs.machine_type }}_run_models_gpu_${{ matrix.folders }}_test_reports tests/${{ matrix.folders }}
 
       - name: Failure short reports
         if: ${{ failure() }}
diff --git a/.github/workflows/push-important-models.yml b/.github/workflows/push-important-models.yml
index ef96539636..82e4d6476a 100644
--- a/.github/workflows/push-important-models.yml
+++ b/.github/workflows/push-important-models.yml
@@ -5,7 +5,6 @@ on:
     branches: [ main ]
 
 env:
-  IS_GITHUB_CI: "1"
   OUTPUT_SLACK_CHANNEL_ID: "C06L2SGMEEA"
   HF_HUB_READ_TOKEN: ${{ secrets.HF_HUB_READ_TOKEN }}
   HF_HOME: /mnt/cache 
@@ -86,7 +85,7 @@ jobs:
       - name: Run FA2 tests
         id: run_fa2_tests
         run:
-          pytest -rs -m "flash_attn_test" --make-reports=${{ matrix.model-name }}_fa2_tests/ tests/${{ matrix.model-name }}/test_modeling_*
+          pytest -rsfE -m "flash_attn_test" --make-reports=${{ matrix.model-name }}_fa2_tests/ tests/${{ matrix.model-name }}/test_modeling_*
       
       - name: "Test suite reports artifacts: ${{ matrix.model-name }}_fa2_tests"
         if: ${{ always() }}
@@ -108,7 +107,7 @@ jobs:
         id: run_integration_tests
         if: always()
         run:
-          pytest -rs -k "IntegrationTest"  --make-reports=tests_integration_${{ matrix.model-name }} tests/${{ matrix.model-name }}/test_modeling_*
+          pytest -rsfE -k "IntegrationTest"  --make-reports=tests_integration_${{ matrix.model-name }} tests/${{ matrix.model-name }}/test_modeling_*
       
       - name: "Test suite reports artifacts: tests_integration_${{ matrix.model-name }}"
         if: ${{ always() }}
diff --git a/.github/workflows/self-pr-slow-ci.yml b/.github/workflows/self-pr-slow-ci.yml
index 10a2156f21..2729c436b1 100644
--- a/.github/workflows/self-pr-slow-ci.yml
+++ b/.github/workflows/self-pr-slow-ci.yml
@@ -110,7 +110,7 @@ jobs:
 
       - name: Run all tests on GPU
         working-directory: /transformers
-        run: python3 -m pytest -v -rs --make-reports=${{ matrix.machine_type }}_run_models_gpu_${{ matrix.folders }}_test_reports tests/${{ matrix.folders }}
+        run: python3 -m pytest -v -rsfE --make-reports=${{ matrix.machine_type }}_run_models_gpu_${{ matrix.folders }}_test_reports tests/${{ matrix.folders }}
 
       - name: Failure short reports
         if: ${{ failure() }}
diff --git a/.github/workflows/ssh-runner.yml b/.github/workflows/ssh-runner.yml
index 3319be368a..7b47c0f437 100644
--- a/.github/workflows/ssh-runner.yml
+++ b/.github/workflows/ssh-runner.yml
@@ -14,7 +14,6 @@ on:
         required: true
 
 env:
-  IS_GITHUB_CI: "1"
   HF_HUB_READ_TOKEN: ${{ secrets.HF_HUB_READ_TOKEN }}
   HF_HOME: /mnt/cache 
   TRANSFORMERS_IS_CI: yes 
diff --git a/docs/source/de/peft.md b/docs/source/de/peft.md
index bdc0684d79..eda8ce9435 100644
--- a/docs/source/de/peft.md
+++ b/docs/source/de/peft.md
@@ -86,10 +86,10 @@ model.load_adapter(peft_model_id)
 Die `bitsandbytes`-Integration unterstützt Datentypen mit 8bit und 4bit Genauigkeit, was für das Laden großer Modelle nützlich ist, weil es Speicher spart (lesen Sie den `bitsandbytes`-Integrations [guide](./quantization#bitsandbytes-integration), um mehr zu erfahren). Fügen Sie die Parameter `load_in_8bit` oder `load_in_4bit` zu [`~PreTrainedModel.from_pretrained`] hinzu und setzen Sie `device_map="auto"`, um das Modell effektiv auf Ihre Hardware zu verteilen:
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 peft_model_id = "ybelkada/opt-350m-lora"
-model = AutoModelForCausalLM.from_pretrained(peft_model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(peft_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 ## Einen neuen Adapter hinzufügen
diff --git a/docs/source/en/agents.md b/docs/source/en/agents.md
index ae9e5db2b7..2cacaed590 100644
--- a/docs/source/en/agents.md
+++ b/docs/source/en/agents.md
@@ -28,8 +28,8 @@ An agent is a system that uses an LLM as its engine, and it has access to functi
 These *tools* are functions for performing a task, and they contain all necessary description for the agent to properly use them.
 
 The agent can be programmed to:
-- devise a series of actions/tools and run them all at once like the `CodeAgent` for example
-- plan and execute actions/tools one by one and wait for the outcome of each action before launching the next one like the `ReactJsonAgent` for example
+- devise a series of actions/tools and run them all at once like the [`CodeAgent`] for example
+- plan and execute actions/tools one by one and wait for the outcome of each action before launching the next one like the [`ReactJsonAgent`] for example
 
 ### Types of agents
 
@@ -42,8 +42,8 @@ This agent has a planning step, then generates python code to execute all its ac
 This is the go-to agent to solve reasoning tasks, since the ReAct framework ([Yao et al., 2022](https://huggingface.co/papers/2210.03629)) makes it really efficient to think on the basis of its previous observations.
 
 We implement two versions of ReactJsonAgent: 
-- [`~ReactJsonAgent`] generates tool calls as a JSON in its output.
-- [`~ReactCodeAgent`] is a new type of ReactJsonAgent that generates its tool calls as blobs of code, which works really well for LLMs that have strong coding performance.
+- [`ReactJsonAgent`] generates tool calls as a JSON in its output.
+- [`ReactCodeAgent`] is a new type of ReactJsonAgent that generates its tool calls as blobs of code, which works really well for LLMs that have strong coding performance.
 
 > [!TIP]
 > Read [Open-source LLMs as LangChain Agents](https://huggingface.co/blog/open-source-llms-as-agents) blog post to learn more the ReAct agent.
@@ -124,7 +124,7 @@ You could use any `llm_engine` method as long as:
 
 You also need a `tools` argument which accepts a list of `Tools`. You can provide an empty list for `tools`, but use the default toolbox with the optional argument `add_base_tools=True`.
 
-Now you can create an agent, like `CodeAgent`, and run it. For convenience, we also provide the `HfEngine` class that uses `huggingface_hub.InferenceClient` under the hood.
+Now you can create an agent, like [`CodeAgent`], and run it. For convenience, we also provide the [`HfEngine`] class that uses `huggingface_hub.InferenceClient` under the hood.
 
 ```python
 from transformers import CodeAgent, HfEngine
@@ -139,7 +139,7 @@ agent.run(
 ```
 
 This will be handy in case of emergency baguette need!
-You can even leave the argument `llm_engine` undefined, and an [~HfEngine] will be created by default.
+You can even leave the argument `llm_engine` undefined, and an [`HfEngine`] will be created by default.
 
 ```python
 from transformers import CodeAgent
@@ -181,13 +181,27 @@ You can also run an agent consecutively for different tasks: each time the attri
 A Python interpreter executes the code on a set of inputs passed along with your tools.
 This should be safe because the only functions that can be called are the tools you provided (especially if it's only tools by Hugging Face) and the print function, so you're already limited in what can be executed.
 
-The Python interpreter also doesn't allow any attribute lookup or imports (which shouldn't be needed for passing inputs/outputs to a small set of functions) so all the most obvious attacks shouldn't be an issue.
+The Python interpreter also doesn't allow imports by default outside of a safe list, so all the most obvious attacks shouldn't be an issue.
+You can still authorize additional imports by passing the authorized modules as a list of strings in argument `additional_authorized_imports` upon initialization of your [`ReactCodeAgent`] or [`CodeAgent`]:
+
+```py
+>>> from transformers import ReactCodeAgent
+
+>>> agent = ReactCodeAgent(tools=[], additional_authorized_imports=['requests', 'bs4'])
+>>>agent.run("Could you get me the title of the page at url 'https://huggingface.co/blog'?")
+
+(...)
+'Hugging Face – Blog'
+```
 
 The execution will stop at any code trying to perform an illegal operation or if there is a regular Python error with the code generated by the agent.
 
+> [!WARNING]
+> The LLM can generate arbitrary code that will then be executed: do not add any unsafe imports!
+
 ### The system prompt
 
-An agent, or rather the LLM that drives the agent, generates an output based on the system prompt. The system prompt can be customized and tailored to the intended task. For example, check the system prompt for the `ReactCodeAgent` (below version is slightly simplified).
+An agent, or rather the LLM that drives the agent, generates an output based on the system prompt. The system prompt can be customized and tailored to the intended task. For example, check the system prompt for the [`ReactCodeAgent`] (below version is slightly simplified).
 
 ```text
 You will be given a task to solve as best you can.
@@ -246,7 +260,7 @@ of the available tools.
 
 A tool is an atomic function to be used by an agent.
 
-You can for instance check the [~PythonInterpreterTool]: it has a name, a description, input descriptions, an output type, and a `__call__` method to perform the action.
+You can for instance check the [`PythonInterpreterTool`]: it has a name, a description, input descriptions, an output type, and a `__call__` method to perform the action.
 
 When the agent is initialized, the tool attributes are used to generate a tool description which is baked into the agent's system prompt. This lets the agent know which tools it can use and why.
 
@@ -259,7 +273,7 @@ Transformers comes with a default toolbox for empowering agents, that you can ad
 - **Speech to text**: given an audio recording of a person talking, transcribe the speech into text ([Whisper](./model_doc/whisper))
 - **Text to speech**: convert text to speech ([SpeechT5](./model_doc/speecht5))
 - **Translation**: translates a given sentence from source language to target language.
-- **Python code interpreter**: runs your the LLM generated Python code in a secure environment. This tool will only be added to [~ReactJsonAgent] if you use `add_base_tools=True`, since code-based tools can already execute Python code
+- **Python code interpreter**: runs your the LLM generated Python code in a secure environment. This tool will only be added to [`ReactJsonAgent`] if you use `add_base_tools=True`, since code-based tools can already execute Python code
 
 
 You can manually use a tool by calling the [`load_tool`] function and a task to perform.
diff --git a/docs/source/en/model_doc/sam.md b/docs/source/en/model_doc/sam.md
index 2fc06193a7..12a87eb5bc 100644
--- a/docs/source/en/model_doc/sam.md
+++ b/docs/source/en/model_doc/sam.md
@@ -81,10 +81,10 @@ processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
 img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
 raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
 mask_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
-segmentation_map = Image.open(requests.get(mask_url, stream=True).raw).convert("RGB")
+segmentation_map = Image.open(requests.get(mask_url, stream=True).raw).convert("1")
 input_points = [[[450, 600]]]  # 2D location of a window in the image
 
-inputs = processor(raw_image, input_points=input_points, segmentation_maps=mask, return_tensors="pt").to(device)
+inputs = processor(raw_image, input_points=input_points, segmentation_maps=segmentation_map, return_tensors="pt").to(device)
 with torch.no_grad():
     outputs = model(**inputs)
 
diff --git a/docs/source/en/peft.md b/docs/source/en/peft.md
index 9e2ac805b2..01fe2fad87 100644
--- a/docs/source/en/peft.md
+++ b/docs/source/en/peft.md
@@ -88,10 +88,10 @@ Check out the [API documentation](#transformers.integrations.PeftAdapterMixin) s
 The `bitsandbytes` integration supports 8bit and 4bit precision data types, which are useful for loading large models because it saves memory (see the `bitsandbytes` integration [guide](./quantization#bitsandbytes-integration) to learn more). Add the `load_in_8bit` or `load_in_4bit` parameters to [`~PreTrainedModel.from_pretrained`] and set `device_map="auto"` to effectively distribute the model to your hardware:
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 peft_model_id = "ybelkada/opt-350m-lora"
-model = AutoModelForCausalLM.from_pretrained(peft_model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(peft_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 ## Add a new adapter
diff --git a/docs/source/en/perf_infer_gpu_one.md b/docs/source/en/perf_infer_gpu_one.md
index 41a5d09a0d..a81fc49381 100644
--- a/docs/source/en/perf_infer_gpu_one.md
+++ b/docs/source/en/perf_infer_gpu_one.md
@@ -354,20 +354,20 @@ If you're curious and interested in learning more about the concepts underlying
 To load a model in 8-bit for inference, use the `load_in_8bit` parameter. The `device_map` parameter is optional, but we recommend setting it to `"auto"` to allow 🤗 Accelerate to automatically and efficiently allocate the model given the available resources in the environment:
 
 ```py
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 If you're loading a model in 8-bit for text generation, you should use the [`~transformers.GenerationMixin.generate`] method instead of the [`Pipeline`] function which is not optimized for 8-bit models and will be slower. Some sampling strategies, like nucleus sampling, are also not supported by the [`Pipeline`] for 8-bit models. You should also place all inputs on the same device as the model:
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 
 prompt = "Hello, my llama is cute"
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
diff --git a/docs/source/it/perf_infer_gpu_one.md b/docs/source/it/perf_infer_gpu_one.md
index 16f77b3b1f..e618ec34a1 100644
--- a/docs/source/it/perf_infer_gpu_one.md
+++ b/docs/source/it/perf_infer_gpu_one.md
@@ -55,10 +55,10 @@ Di seguito sono riportate alcune note per aiutarvi a utilizzare questo modulo, o
 Dopo aver installato le librerie necessarie, per caricare il tuo modello mixed 8-bit è il seguente:
 
 ```py
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 Per la generazione di testo, si consiglia di:
@@ -69,11 +69,11 @@ Per la generazione di testo, si consiglia di:
 Ecco un semplice esempio:
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 
 text = "Hello, my llama is cute"
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -87,7 +87,7 @@ outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
 Usare il seguente modo caricare il modello mixed-8bit su più GPU (stesso comando della configurazione a GPU singola):
 ```py
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 Puoi controllare la RAM della GPU che si vuole allocare su ogni GPU usando `accelerate`. Utilizzare l'argomento `max_memory` come segue:
 
diff --git a/docs/source/ja/main_classes/quantization.md b/docs/source/ja/main_classes/quantization.md
index 3af3130a84..a93d06b257 100644
--- a/docs/source/ja/main_classes/quantization.md
+++ b/docs/source/ja/main_classes/quantization.md
@@ -245,12 +245,12 @@ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_i
 
 ```python
 # pip install transformers accelerate bitsandbytes
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_id = "bigscience/bloom-1b7"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 次に、通常 [`PreTrainedModel`] を使用するのと同じようにモデルを使用します。
@@ -321,9 +321,9 @@ model_double_quant = AutoModelForCausalLM.from_pretrained(model_id, quantization
 この機能を使用できるようにするには、必ず `bitsandbytes>0.37.2` を使用してください (この記事の執筆時点では、`bitsandbytes==0.38.0.post1` でテストしました)。
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
 
 model.push_to_hub("bloom-560m-8bit")
diff --git a/docs/source/ja/peft.md b/docs/source/ja/peft.md
index 5cc687f70b..c3d195adbd 100644
--- a/docs/source/ja/peft.md
+++ b/docs/source/ja/peft.md
@@ -91,10 +91,10 @@ model.load_adapter(peft_model_id)
 `bitsandbytes` 統合は、8ビットおよび4ビットの精度データ型をサポートしており、大規模なモデルを読み込む際にメモリを節約するのに役立ちます（詳細については `bitsandbytes` 統合の[ガイド](./quantization#bitsandbytes-integration)を参照してください）。[`~PreTrainedModel.from_pretrained`] に `load_in_8bit` または `load_in_4bit` パラメータを追加し、`device_map="auto"` を設定してモデルを効果的にハードウェアに分散配置できます：
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 peft_model_id = "ybelkada/opt-350m-lora"
-model = AutoModelForCausalLM.from_pretrained(peft_model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(peft_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 ## Add a new adapter
diff --git a/docs/source/ja/perf_infer_gpu_one.md b/docs/source/ja/perf_infer_gpu_one.md
index 6d7466e022..d6a9b30916 100644
--- a/docs/source/ja/perf_infer_gpu_one.md
+++ b/docs/source/ja/perf_infer_gpu_one.md
@@ -357,10 +357,10 @@ Int8混合精度行列分解は、行列乗算を2つのストリームに分割
 必要なライブラリをインストールした後、ミックス 8 ビットモデルを読み込む方法は次の通りです：
 
 ```py
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 以下はシンプルな例です：
@@ -370,11 +370,11 @@ model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto",
 
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 
 prompt = "Hello, my llama is cute"
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -388,7 +388,7 @@ outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
 
 ```py
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 `accelerate`を使用して各GPUに割り当てるGPU RAMを制御する際には、以下のように`max_memory`引数を使用します：
diff --git a/docs/source/ko/peft.md b/docs/source/ko/peft.md
index 90327e62c2..d4ef0ba539 100644
--- a/docs/source/ko/peft.md
+++ b/docs/source/ko/peft.md
@@ -86,10 +86,10 @@ model.load_adapter(peft_model_id)
 `bitsandbytes` 통합은 8비트와 4비트 정밀도 데이터 유형을 지원하므로 큰 모델을 가져올 때 유용하면서 메모리도 절약합니다. 모델을 하드웨어에 효과적으로 분배하려면 [`~PreTrainedModel.from_pretrained`]에 `load_in_8bit` 또는 `load_in_4bit` 매개변수를 추가하고 `device_map="auto"`를 설정하세요:
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 peft_model_id = "ybelkada/opt-350m-lora"
-model = AutoModelForCausalLM.from_pretrained(peft_model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(peft_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 ## 새 어댑터 추가 [[add-a-new-adapter]]
diff --git a/docs/source/ko/perf_infer_gpu_one.md b/docs/source/ko/perf_infer_gpu_one.md
index 73cef858b9..d6ddca6cd0 100644
--- a/docs/source/ko/perf_infer_gpu_one.md
+++ b/docs/source/ko/perf_infer_gpu_one.md
@@ -127,10 +127,10 @@ Int8 혼합 정밀도 행렬 분해는 행렬 곱셈을 두 개의 스트림으
 필요한 라이브러리를 설치한 후 혼합 8비트 모델을 가져오는 방법은 다음과 같습니다:
 
 ```py
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 텍스트 생성의 경우:
@@ -141,11 +141,11 @@ model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto",
 다음은 간단한 예입니다:
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_name = "bigscience/bloom-2b5"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 
 prompt = "Hello, my llama is cute"
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -159,7 +159,7 @@ outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
 다중 GPU에서 혼합 8비트 모델을 로드하는 방법은 단일 GPU 설정과 동일합니다(동일한 명령어 사용):
 ```py
 model_name = "bigscience/bloom-2b5"
-model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 하지만 `accelerate`를 사용하여 각 GPU에 할당할 GPU RAM을 제어할 수 있습니다. 다음과 같이 `max_memory` 인수를 사용하세요:
 
diff --git a/docs/source/zh/main_classes/quantization.md b/docs/source/zh/main_classes/quantization.md
index 3c7e4d9212..d303906a99 100644
--- a/docs/source/zh/main_classes/quantization.md
+++ b/docs/source/zh/main_classes/quantization.md
@@ -360,12 +360,12 @@ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_i
 
 ```python
 # pip install transformers accelerate bitsandbytes
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_id = "bigscience/bloom-1b7"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 然后，像通常使用 `PreTrainedModel` 一样使用您的模型。
@@ -441,9 +441,9 @@ model_double_quant = AutoModelForCausalLM.from_pretrained(model_id, quantization
 
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
 
 model.push_to_hub("bloom-560m-8bit")
diff --git a/docs/source/zh/peft.md b/docs/source/zh/peft.md
index 4241a15c00..de7ae6d155 100644
--- a/docs/source/zh/peft.md
+++ b/docs/source/zh/peft.md
@@ -86,10 +86,10 @@ model.load_adapter(peft_model_id)
 `bitsandbytes`集成支持8bit和4bit精度数据类型，这对于加载大模型非常有用，因为它可以节省内存（请参阅`bitsandbytes`[指南](./quantization#bitsandbytes-integration)以了解更多信息）。要有效地将模型分配到您的硬件，请在[`~PreTrainedModel.from_pretrained`]中添加`load_in_8bit`或`load_in_4bit`参数，并将`device_map="auto"`设置为：
 
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 peft_model_id = "ybelkada/opt-350m-lora"
-model = AutoModelForCausalLM.from_pretrained(peft_model_id, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(peft_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True))
 ```
 
 ## 添加新的adapter
diff --git a/src/transformers/agents/agents.py b/src/transformers/agents/agents.py
index 64e810eb91..ad0b9fecc3 100644
--- a/src/transformers/agents/agents.py
+++ b/src/transformers/agents/agents.py
@@ -26,7 +26,7 @@ from .agent_types import AgentAudio, AgentImage, AgentText
 from .default_tools import BASE_PYTHON_TOOLS, FinalAnswerTool, setup_default_tools
 from .llm_engine import HfEngine, MessageRole
 from .prompts import DEFAULT_CODE_SYSTEM_PROMPT, DEFAULT_REACT_CODE_SYSTEM_PROMPT, DEFAULT_REACT_JSON_SYSTEM_PROMPT
-from .python_interpreter import evaluate_python_code
+from .python_interpreter import LIST_SAFE_MODULES, evaluate_python_code
 from .tools import (
     DEFAULT_TOOL_DESCRIPTION_TEMPLATE,
     Tool,
@@ -84,8 +84,14 @@ def parse_json_blob(json_blob: str) -> Dict[str, str]:
         return json_data
     except json.JSONDecodeError as e:
         place = e.pos
+        if json_blob[place - 1 : place + 2] == "},\n":
+            raise ValueError(
+                "JSON is invalid: you probably tried to provide multiple tool calls in one action. PROVIDE ONLY ONE TOOL CALL."
+            )
         raise ValueError(
-            f"The JSON blob you used is invalid: due to the following error: {e}. JSON blob was: {json_blob}, decoding failed at '{json_blob[place-4:place+5]}'."
+            f"The JSON blob you used is invalid due to the following error: {e}.\n"
+            f"JSON blob was: {json_blob}, decoding failed on that specific part of the blob:\n"
+            f"'{json_blob[place-4:place+5]}'."
         )
     except Exception as e:
         raise ValueError(f"Error in parsing the JSON blob: {e}")
@@ -347,6 +353,7 @@ class Agent:
         return self._toolbox
 
     def initialize_for_run(self, task: str, **kwargs):
+        self.token_count = 0
         self.task = task
         if len(kwargs) > 0:
             self.task += f"\nYou have been provided with these initial arguments: {str(kwargs)}."
@@ -380,7 +387,7 @@ class Agent:
                 message_content = (
                     "Error: "
                     + str(step_log["error"])
-                    + "\nNow let's retry: take care not to repeat previous errors! Try to adopt different approaches.\n"
+                    + "\nNow let's retry: take care not to repeat previous errors! If you have retried several times, try a completely different approach.\n"
                 )
             elif "observation" in step_log:
                 message_content = f"Observation: {step_log['observation']}"
@@ -409,6 +416,9 @@ class Agent:
                 )
         return memory
 
+    def get_succinct_logs(self):
+        return [{key: value for key, value in log.items() if key != "agent_memory"} for log in self.logs]
+
     def extract_action(self, llm_output: str, split_token: str) -> str:
         """
         Parse action from the LLM output
@@ -486,6 +496,7 @@ class CodeAgent(Agent):
         llm_engine: Callable = HfEngine(),
         system_prompt: str = DEFAULT_CODE_SYSTEM_PROMPT,
         tool_description_template: str = DEFAULT_TOOL_DESCRIPTION_TEMPLATE,
+        additional_authorized_imports: List[str] = [],
         **kwargs,
     ):
         super().__init__(
@@ -504,6 +515,7 @@ class CodeAgent(Agent):
             )
 
         self.python_evaluator = evaluate_python_code
+        self.additional_authorized_imports = additional_authorized_imports
 
     def parse_code_blob(self, result: str) -> str:
         """
@@ -544,7 +556,7 @@ class CodeAgent(Agent):
         self.prompt = [prompt_message, task_message]
         self.logger.info("====Executing with this prompt====")
         self.logger.info(self.prompt)
-        llm_output = self.llm_engine(self.prompt, stop_sequences=["<end_code>"])
+        llm_output = self.llm_engine(self.prompt, stop_sequences=["<end_action>"])
 
         if return_generated_code:
             return llm_output
@@ -563,7 +575,12 @@ class CodeAgent(Agent):
         self.log_code_action(code_action)
         try:
             available_tools = {**BASE_PYTHON_TOOLS.copy(), **self.toolbox.tools}
-            output = self.python_evaluator(code_action, available_tools, state=self.state)
+            output = self.python_evaluator(
+                code_action,
+                available_tools,
+                state=self.state,
+                authorized_imports=LIST_SAFE_MODULES + self.additional_authorized_imports,
+            )
             self.logger.info(self.state["print_outputs"])
             return output
         except Exception as e:
@@ -597,7 +614,29 @@ class ReactAgent(Agent):
         if "final_answer" not in self._toolbox.tools:
             self._toolbox.add_tool(FinalAnswerTool())
 
-    def run(self, task: str, **kwargs):
+    def provide_final_answer(self, task) -> str:
+        """
+        This method provides a final answer to the task, based on the logs of the agent's interactions.
+        """
+        self.prompt = [
+            {
+                "role": MessageRole.SYSTEM,
+                "content": "An agent tried to answer an user query but it got stuck and failed to do so. You are tasked with providing an answer instead. Here is the agent's memory:",
+            }
+        ]
+        self.prompt += self.write_inner_memory_from_logs()[1:]
+        self.prompt += [
+            {
+                "role": MessageRole.USER,
+                "content": f"Based on the above, please provide an answer to the following user request:\n{task}",
+            }
+        ]
+        try:
+            return self.llm_engine(self.prompt)
+        except Exception as e:
+            return f"Error in generating final llm output: {e}."
+
+    def run(self, task: str, stream: bool = False, **kwargs):
         """
         Runs the agent for the given task.
 
@@ -614,41 +653,62 @@ class ReactAgent(Agent):
         agent.run("What is the result of 2 power 3.7384?")
         ```
         """
+        if stream:
+            return self.stream_run(task, **kwargs)
+        else:
+            return self.direct_run(task, **kwargs)
+
+    def stream_run(self, task: str, **kwargs):
         self.initialize_for_run(task, **kwargs)
 
         final_answer = None
         iteration = 0
         while final_answer is None and iteration < self.max_iterations:
             try:
-                final_answer = self.step()
+                step_logs = self.step()
+                if "final_answer" in step_logs:
+                    final_answer = step_logs["final_answer"]
             except AgentError as e:
                 self.logger.error(e, exc_info=1)
                 self.logs[-1]["error"] = e
             finally:
                 iteration += 1
+                yield self.logs[-1]
 
         if final_answer is None and iteration == self.max_iterations:
             error_message = "Reached max iterations."
-            self.logs.append({"error": AgentMaxIterationsError(error_message)})
+            final_step_log = {"error": AgentMaxIterationsError(error_message)}
+            self.logs.append(final_step_log)
             self.logger.error(error_message, exc_info=1)
+            final_answer = self.provide_final_answer(task)
+            final_step_log["final_answer"] = final_answer
+            yield final_step_log
+
+        yield final_answer
 
-            self.prompt = [
-                {
-                    "role": MessageRole.SYSTEM,
-                    "content": "An agent tried to answer a user query but it failed to do so. You are tasked with providing an answer instead. Here is the agent's memory:",
-                }
-            ]
-            self.prompt += self.write_inner_memory_from_logs()[1:]
-            self.prompt += [
-                {
-                    "role": MessageRole.USER,
-                    "content": f"Based on the above, please provide an answer to the following user request:\n{task}",
-                }
-            ]
+    def direct_run(self, task: str, **kwargs):
+        self.initialize_for_run(task, **kwargs)
+
+        final_answer = None
+        iteration = 0
+        while final_answer is None and iteration < self.max_iterations:
             try:
-                final_answer = self.llm_engine(self.prompt, stop_sequences=["Observation:"])
-            except Exception as e:
-                final_answer = f"Error in generating final llm output: {e}."
+                step_logs = self.step()
+                if "final_answer" in step_logs:
+                    final_answer = step_logs["final_answer"]
+            except AgentError as e:
+                self.logger.error(e, exc_info=1)
+                self.logs[-1]["error"] = e
+            finally:
+                iteration += 1
+
+        if final_answer is None and iteration == self.max_iterations:
+            error_message = "Reached max iterations."
+            final_step_log = {"error": AgentMaxIterationsError(error_message)}
+            self.logs.append(final_step_log)
+            self.logger.error(error_message, exc_info=1)
+            final_answer = self.provide_final_answer(task)
+            final_step_log["final_answer"] = final_answer
 
         return final_answer
 
@@ -683,22 +743,24 @@ class ReactJsonAgent(ReactAgent):
         """
         agent_memory = self.write_inner_memory_from_logs()
 
-        self.logs[-1]["agent_memory"] = agent_memory.copy()
         self.prompt = agent_memory
         self.logger.debug("===== New step =====")
 
         # Add new step in logs
-        self.logs.append({})
+        current_step_logs = {}
+        self.logs.append(current_step_logs)
+        current_step_logs["agent_memory"] = agent_memory.copy()
+
         self.logger.info("===== Calling LLM with this last message: =====")
         self.logger.info(self.prompt[-1])
 
         try:
-            llm_output = self.llm_engine(self.prompt, stop_sequences=["Observation:"])
+            llm_output = self.llm_engine(self.prompt, stop_sequences=["<end_action>", "Observation:"])
         except Exception as e:
             raise AgentGenerationError(f"Error in generating llm output: {e}.")
         self.logger.debug("===== Output message of the LLM: =====")
         self.logger.debug(llm_output)
-        self.logs[-1]["llm_output"] = llm_output
+        current_step_logs["llm_output"] = llm_output
 
         # Parse
         self.logger.debug("===== Extracting action =====")
@@ -709,8 +771,8 @@ class ReactJsonAgent(ReactAgent):
         except Exception as e:
             raise AgentParsingError(f"Could not parse the given action: {e}.")
 
-        self.logs[-1]["rationale"] = rationale
-        self.logs[-1]["tool_call"] = {"tool_name": tool_name, "tool_arguments": arguments}
+        current_step_logs["rationale"] = rationale
+        current_step_logs["tool_call"] = {"tool_name": tool_name, "tool_arguments": arguments}
 
         # Execute
         self.logger.warning(f"Calling tool: '{tool_name}' with arguments: {arguments}")
@@ -721,7 +783,8 @@ class ReactJsonAgent(ReactAgent):
                 answer = arguments
             if answer in self.state:  # if the answer is a state variable, return the value
                 answer = self.state[answer]
-            return answer
+            current_step_logs["final_answer"] = answer
+            return current_step_logs
         else:
             observation = self.execute_tool_call(tool_name, arguments)
             observation_type = type(observation)
@@ -740,8 +803,8 @@ class ReactJsonAgent(ReactAgent):
                 updated_information = f"Stored '{observation_name}' in memory."
 
             self.logger.info(updated_information)
-            self.logs[-1]["observation"] = updated_information
-            return None
+            current_step_logs["observation"] = updated_information
+            return current_step_logs
 
 
 class ReactCodeAgent(ReactAgent):
@@ -757,6 +820,7 @@ class ReactCodeAgent(ReactAgent):
         llm_engine: Callable = HfEngine(),
         system_prompt: str = DEFAULT_REACT_CODE_SYSTEM_PROMPT,
         tool_description_template: str = DEFAULT_TOOL_DESCRIPTION_TEMPLATE,
+        additional_authorized_imports: List[str] = [],
         **kwargs,
     ):
         super().__init__(
@@ -775,6 +839,7 @@ class ReactCodeAgent(ReactAgent):
             )
 
         self.python_evaluator = evaluate_python_code
+        self.additional_authorized_imports = additional_authorized_imports
 
     def step(self):
         """
@@ -782,26 +847,27 @@ class ReactCodeAgent(ReactAgent):
         The errors are raised here, they are caught and logged in the run() method.
         """
         agent_memory = self.write_inner_memory_from_logs()
-        self.logs[-1]["agent_memory"] = agent_memory.copy()
 
         self.prompt = agent_memory.copy()
 
         self.logger.debug("===== New step =====")
 
         # Add new step in logs
-        self.logs.append({})
+        current_step_logs = {}
+        self.logs.append(current_step_logs)
+        current_step_logs["agent_memory"] = agent_memory.copy()
 
         self.logger.info("===== Calling LLM with these last messages: =====")
         self.logger.info(self.prompt[-2:])
 
         try:
-            llm_output = self.llm_engine(self.prompt, stop_sequences=["<end_code>", "Observation:"])
+            llm_output = self.llm_engine(self.prompt, stop_sequences=["<end_action>", "Observation:"])
         except Exception as e:
             raise AgentGenerationError(f"Error in generating llm output: {e}.")
 
         self.logger.debug("===== Output message of the LLM: =====")
         self.logger.debug(llm_output)
-        self.logs[-1]["llm_output"] = llm_output
+        current_step_logs["llm_output"] = llm_output
 
         # Parse
         self.logger.debug("===== Extracting action =====")
@@ -813,18 +879,23 @@ class ReactCodeAgent(ReactAgent):
             error_msg = f"Error in code parsing: {e}. Make sure to provide correct code"
             raise AgentParsingError(error_msg)
 
-        self.logs[-1]["rationale"] = rationale
-        self.logs[-1]["tool_call"] = {"tool_name": "code interpreter", "tool_arguments": code_action}
+        current_step_logs["rationale"] = rationale
+        current_step_logs["tool_call"] = {"tool_name": "code interpreter", "tool_arguments": code_action}
 
         # Execute
         self.log_code_action(code_action)
         try:
             available_tools = {**BASE_PYTHON_TOOLS.copy(), **self.toolbox.tools}
-            result = self.python_evaluator(code_action, available_tools, state=self.state)
+            result = self.python_evaluator(
+                code_action,
+                available_tools,
+                state=self.state,
+                authorized_imports=LIST_SAFE_MODULES + self.additional_authorized_imports,
+            )
             information = self.state["print_outputs"]
             self.logger.warning("Print outputs:")
             self.logger.log(32, information)
-            self.logs[-1]["observation"] = information
+            current_step_logs["observation"] = information
         except Exception as e:
             error_msg = f"Failed while trying to execute the code below:\n{CustomFormatter.reset + code_action + CustomFormatter.reset}\nThis failed due to the following error:\n{str(e)}"
             if "'dict' object has no attribute 'read'" in str(e):
@@ -834,5 +905,5 @@ class ReactCodeAgent(ReactAgent):
             if line[: len("final_answer")] == "final_answer":
                 self.logger.warning(">>> Final answer:")
                 self.logger.log(32, result)
-                return result
-        return None
+                current_step_logs["final_answer"] = result
+        return current_step_logs
diff --git a/src/transformers/agents/llm_engine.py b/src/transformers/agents/llm_engine.py
index b696084090..76458b0267 100644
--- a/src/transformers/agents/llm_engine.py
+++ b/src/transformers/agents/llm_engine.py
@@ -61,7 +61,6 @@ def get_clean_message_list(message_list: List[Dict[str, str]], role_conversions:
 
 
 llama_role_conversions = {
-    MessageRole.SYSTEM: MessageRole.USER,
     MessageRole.TOOL_RESPONSE: MessageRole.USER,
 }
 
@@ -72,20 +71,14 @@ class HfEngine:
         self.client = InferenceClient(model=self.model, timeout=120)
 
     def __call__(self, messages: List[Dict[str, str]], stop_sequences=[]) -> str:
-        if "Meta-Llama-3" in self.model:
-            if "
" not in stop_sequences:
-                stop_sequences.append("
")
-            if "!!!!!" not in stop_sequences:
-                stop_sequences.append("!!!!!")
-
         # Get clean message list
         messages = get_clean_message_list(messages, role_conversions=llama_role_conversions)
 
-        # Get answer
+        # Get LLM output
         response = self.client.chat_completion(messages, stop=stop_sequences, max_tokens=1500)
         response = response.choices[0].message.content
 
-        # Remove stop sequences from the answer
+        # Remove stop sequences from LLM output
         for stop_seq in stop_sequences:
             if response[-len(stop_seq) :] == stop_seq:
                 response = response[: -len(stop_seq)]
diff --git a/src/transformers/agents/prompts.py b/src/transformers/agents/prompts.py
index 80c65a5144..f76734b09f 100644
--- a/src/transformers/agents/prompts.py
+++ b/src/transformers/agents/prompts.py
@@ -68,7 +68,7 @@ translated_question = translator(question=question, src_lang="French", tgt_lang=
 print(f"The translated question is {translated_question}.")
 answer = image_qa(image=image, question=translated_question)
 print(f"The answer is {answer}")
-```<end_code>
+```<end_action>
 
 ---
 Task: "Identify the oldest person in the `document` and create an image showcasing the result."
@@ -79,7 +79,7 @@ Code:
 answer = document_qa(document, question="What is the oldest person?")
 print(f"The answer is {answer}.")
 image = image_generator(answer)
-```<end_code>
+```<end_action>
 
 ---
 Task: "Generate an image using the text given in the variable `caption`."
@@ -88,7 +88,7 @@ I will use the following tool: `image_generator` to generate an image.
 Code:
 ```py
 image = image_generator(prompt=caption)
-```<end_code>
+```<end_action>
 
 ---
 Task: "Summarize the text given in the variable `text` and read it out loud."
@@ -99,7 +99,7 @@ Code:
 summarized_text = summarizer(text)
 print(f"Summary: {summarized_text}")
 audio_summary = text_reader(summarized_text)
-```<end_code>
+```<end_action>
 
 ---
 Task: "Answer the question in the variable `question` about the text in the variable `text`. Use the answer to generate an image."
@@ -110,7 +110,7 @@ Code:
 answer = text_qa(text=text, question=question)
 print(f"The answer is {answer}.")
 image = image_generator(answer)
-```<end_code>
+```<end_action>
 
 ---
 Task: "Caption the following `image`."
@@ -119,39 +119,32 @@ I will use the following tool: `image_captioner` to generate a caption for the i
 Code:
 ```py
 caption = image_captioner(image)
-```<end_code>
+```<end_action>
 
 ---
 Above example were using tools that might not exist for you. You only have acces to those Tools:
 <<tool_names>>
 
 Remember to make sure that variables you use are all defined.
-Be sure to provide a 'Code:\n```' sequence before the code and '```<end_code>' after, else you will get an error.
+Be sure to provide a 'Code:\n```' sequence before the code and '```<end_action>' after, else you will get an error.
 DO NOT pass the arguments as a dict as in 'answer = ask_search_agent({'query': "What is the place where James Bond lives?"})', but use the arguments directly as in 'answer = ask_search_agent(query="What is the place where James Bond lives?")'.
 
-Now Begin!
+Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000.
 """
 
 
-DEFAULT_REACT_JSON_SYSTEM_PROMPT = """You will be given a task to solve as best you can. You have access to the following tools:
-<<tool_descriptions>>
-
-The way you use the tools is by specifying a json blob.
-Specifically, this json should have a `action` key (name of the tool to use) and a `action_input` key (input to the tool).
+DEFAULT_REACT_JSON_SYSTEM_PROMPT = """You will be given a task to solve as best you can. To do so, you have been given access to the following tools: <<tool_names>>
+The way you use the tools is by specifying a json blob, ending with '<end_action>'.
+Specifically, this json should have an `action` key (name of the tool to use) and an `action_input` key (input to the tool).
 
 The $ACTION_JSON_BLOB should only contain a SINGLE action, do NOT return a list of multiple actions. It should be formatted in json. Do not try to escape special characters. Here is the template of a valid $ACTION_JSON_BLOB:
-Action:
 {
   "action": $TOOL_NAME,
   "action_input": $INPUT
-}
+}<end_action>
 
 Make sure to have the $INPUT as a dictionnary in the right format for the tool you are using, and do not put variable names as input if you can find the right values.
 
-You will be given:
-
-Task: the task you are given.
-
 You should ALWAYS use the following format:
 
 Thought: you should always think about one action to take. Then use the action as follows:
@@ -171,14 +164,14 @@ Action:
 {
   "action": "image_transformer",
   "action_input": {"image": "image_1.jpg"}
-}
+}<end_action>
 
 To provide the final answer to the task, use an action blob with "action": "final_answer" tool. It is the only way to complete the task, else you will be stuck on a loop. So your final output should look like this:
 Action:
 {
   "action": "final_answer",
   "action_input": {"answer": "insert your final answer here"}
-}
+}<end_action>
 
 
 Here are a few examples using notional tools:
@@ -190,7 +183,7 @@ Action:
 {
   "action": "document_qa",
   "action_input": {"document": "document.pdf", "question": "Who is the oldest person mentioned?"}
-}
+}<end_action>
 Observation: "The oldest person in the document is John Doe, a 55 year old lumberjack living in Newfoundland."
 
 
@@ -199,7 +192,7 @@ Action:
 {
   "action": "image_generator",
   "action_input": {"text": ""A portrait of John Doe, a 55-year-old man living in Canada.""}
-}
+}<end_action>
 Observation: "image.png"
 
 Thought: I will now return the generated image.
@@ -207,7 +200,7 @@ Action:
 {
   "action": "final_answer",
   "action_input": "image.png"
-}
+}<end_action>
 
 ---
 Task: "What is the result of the following operation: 5 + 3 + 1294.678?"
@@ -217,7 +210,7 @@ Action:
 {
     "action": "python_interpreter",
     "action_input": {"code": "5 + 3 + 1294.678"}
-}
+}<end_action>
 Observation: 1302.678
 
 Thought: Now that I know the result, I will now return it.
@@ -225,7 +218,7 @@ Action:
 {
   "action": "final_answer",
   "action_input": "1302.678"
-}
+}<end_action>
 
 ---
 Task: "Which city has the highest population , Guangzhou or Shanghai?"
@@ -235,7 +228,7 @@ Action:
 {
     "action": "search",
     "action_input": "Population Guangzhou"
-}
+}<end_action>
 Observation: ['Guangzhou has a population of 15 million inhabitants as of 2021.']
 
 
@@ -252,28 +245,30 @@ Action:
 {
   "action": "final_answer",
   "action_input": "Shanghai"
-}
+}<end_action>
 
 
 Above example were using notional tools that might not exist for you. You only have acces to those tools:
-<<tool_names>>
-ALWAYS provide a 'Thought:' and an 'Action:' sequence. You MUST provide at least the 'Action:' sequence to move forward.
+<<tool_descriptions>>
+
+Here are the rules you should always follow to solve your task:
+1. ALWAYS provide a 'Thought:' sequence, and an 'Action:' sequence that ends with <end_action>, else you will fail.
+2. Always use the right arguments for the tools. Never use variable names in the 'action_input' field, use the value instead.
+3. Call a tool only when needed: do not call the search agent if you do not need information, try to solve the task yourself.
+4. Never re-do a tool call that you previously did with the exact same parameters.
 
-Now begin!
+Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000.
 """
 
 
 DEFAULT_REACT_CODE_SYSTEM_PROMPT = """You will be given a task to solve as best you can.
-You have access to the following tools:
-<<tool_descriptions>>
-
+To do so, you have been given access to *tools*: these tools are basically Python functions which you can call with code.
 To solve the task, you must plan forward to proceed in a series of steps, in a cycle of 'Thought:', 'Code:', and 'Observation:' sequences.
 
-At each step, in the 'Thought:' sequence, you should first explain your reasoning towards solving the task, then the tools that you want to use.
-Then in the 'Code:' sequence, you shold write the code in simple Python. The code sequence must end with '/End code' sequence.
+At each step, in the 'Thought:' sequence, you should first explain your reasoning towards solving the task and the tools that you want to use.
+Then in the 'Code:' sequence, you should write the code in simple Python. The code sequence must end with '<end_action>' sequence.
 During each intermediate step, you can use 'print()' to save whatever important information you will then need.
-These print outputs will then be available in the 'Observation:' field, for using this information as input for the next step.
-
+These print outputs will then appear in the 'Observation:' field, which will be available as input for the next step.
 In the end you have to return a final answer using the `final_answer` tool.
 
 Here are a few examples using notional tools:
@@ -285,7 +280,7 @@ Code:
 ```py
 answer = document_qa(document=document, question="Who is the oldest person mentioned?")
 print(answer)
-```<end_code>
+```<end_action>
 Observation: "The oldest person in the document is John Doe, a 55 year old lumberjack living in Newfoundland."
 
 Thought: I will now generate an image showcasing the oldest person.
@@ -294,7 +289,7 @@ Code:
 ```py
 image = image_generator("A portrait of John Doe, a 55-year-old man living in Canada.")
 final_answer(image)
-```<end_code>
+```<end_action>
 
 ---
 Task: "What is the result of the following operation: 5 + 3 + 1294.678?"
@@ -305,10 +300,10 @@ Code:
 ```py
 result = 5 + 3 + 1294.678
 final_answer(result)
-```<end_code>
+```<end_action>
 
 ---
-Task: "Which city has the highest population , Guangzhou or Shanghai?"
+Task: "Which city has the highest population: Guangzhou or Shanghai?"
 
 Thought: I need to get the populations for both cities and compare them: I will use the tool `search` to get the population of both cities.
 Code:
@@ -317,7 +312,7 @@ population_guangzhou = search("Guangzhou population")
 print("Population Guangzhou:", population_guangzhou)
 population_shanghai = search("Shanghai population")
 print("Population Shanghai:", population_shanghai)
-```<end_code>
+```<end_action>
 Observation:
 Population Guangzhou: ['Guangzhou has a population of 15 million inhabitants as of 2021.']
 Population Shanghai: '26 million (2019)'
@@ -326,7 +321,7 @@ Thought: Now I know that Shanghai has the highest population.
 Code:
 ```py
 final_answer("Shanghai")
-```<end_code>
+```<end_action>
 
 ---
 Task: "What is the current age of the pope, raised to the power 0.36?"
@@ -336,7 +331,7 @@ Code:
 ```py
 pope_age = search(query="current pope age")
 print("Pope age:", pope_age)
-```<end_code>
+```<end_action>
 Observation:
 Pope age: "The pope Francis is currently 85 years old."
 
@@ -345,20 +340,21 @@ Code:
 ```py
 pope_current_age = 85 ** 0.36
 final_answer(pope_current_age)
-```<end_code>
-
+```<end_action>
 
 Above example were using notional tools that might not exist for you. You only have acces to those tools:
-<<tool_names>>
-You also can perform computations in the python code you generate.
 
-Always provide a 'Thought:' and a 'Code:\n```py' sequence ending with '```<end_code>' sequence. You MUST provide at least the 'Code:' sequence to move forward.
+<<tool_descriptions>>
 
-Remember to not perform too many operations in a single code block! You should split the task into intermediate code blocks.
-Print results at the end of each step to save the intermediate results. Then use final_answer() to return the final result.
+You also can perform computations in the Python code that you generate.
 
-Remember to make sure that variables you use are all defined.
-DO NOT pass the arguments as a dict as in 'answer = ask_search_agent({'query': "What is the place where James Bond lives?"})', but use the arguments directly as in 'answer = ask_search_agent(query="What is the place where James Bond lives?")'.
+Here are the rules you should always follow to solve your task:
+1. Always provide a 'Thought:' sequence, and a 'Code:\n```py' sequence ending with '```<end_action>' sequence, else you will fail.
+2. Use only variables that you have defined!
+3. Always use the right arguments for the tools. DO NOT pass the arguments as a dict as in 'answer = ask_search_agent({'query': "What is the place where James Bond lives?"})', but use the arguments directly as in 'answer = ask_search_agent(query="What is the place where James Bond lives?")'.
+4. Take care to not chain too many sequential tool calls in the same code block, especially when the output format is unpredictable. For instance, a call to search has an unpredictable return format, so do not have another tool call that depends on its output in the same block: rather output results with print() to use them in the next block.
+5. Call a tool only when needed, and never re-do a tool call that you previously did with the exact same parameters.
+6. Don't name any new variable with the same name as a tool: for instance don't name a variable 'final_answer'.
 
-Now Begin!
+Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000.
 """
diff --git a/src/transformers/agents/python_interpreter.py b/src/transformers/agents/python_interpreter.py
index 8ca1cd1820..992e9d14f1 100644
--- a/src/transformers/agents/python_interpreter.py
+++ b/src/transformers/agents/python_interpreter.py
@@ -15,9 +15,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import ast
+import builtins
 import difflib
 from collections.abc import Mapping
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 
 class InterpretorError(ValueError):
@@ -29,7 +30,25 @@ class InterpretorError(ValueError):
     pass
 
 
-LIST_SAFE_MODULES = ["random", "math", "time", "queue", "itertools", "re", "stat", "statistics", "unicodedata"]
+ERRORS = {
+    name: getattr(builtins, name)
+    for name in dir(builtins)
+    if isinstance(getattr(builtins, name), type) and issubclass(getattr(builtins, name), BaseException)
+}
+
+
+LIST_SAFE_MODULES = [
+    "random",
+    "collections",
+    "math",
+    "time",
+    "queue",
+    "itertools",
+    "re",
+    "stat",
+    "statistics",
+    "unicodedata",
+]
 
 
 class BreakException(Exception):
@@ -87,21 +106,62 @@ def evaluate_while(while_loop, state, tools):
     return None
 
 
-def evaluate_function_def(function_def, state, tools):
-    def create_function(func_def, state, tools):
-        def new_func(*args):
-            new_state = state.copy()
-            for arg, val in zip(func_def.args.args, args):
-                new_state[arg.arg] = val
-            result = None
-            for node in func_def.body:
-                result = evaluate_ast(node, new_state, tools)
-            return result
+def create_function(func_def, state, tools):
+    def new_func(*args, **kwargs):
+        func_state = state.copy()
+        arg_names = [arg.arg for arg in func_def.args.args]
+        for name, value in zip(arg_names, args):
+            func_state[name] = value
+        if func_def.args.vararg:
+            vararg_name = func_def.args.vararg.arg
+            func_state[vararg_name] = args
+        if func_def.args.kwarg:
+            kwarg_name = func_def.args.kwarg.arg
+            func_state[kwarg_name] = kwargs
+
+        # Update function state with self and __class__
+        if func_def.args.args and func_def.args.args[0].arg == "self":
+            if args:
+                func_state["self"] = args[0]
+                func_state["__class__"] = args[0].__class__
+
+        result = None
+        for stmt in func_def.body:
+            result = evaluate_ast(stmt, func_state, tools)
+        return result
 
-        return new_func
+    return new_func
 
-    tools[function_def.name] = create_function(function_def, state, tools)
-    return None
+
+def create_class(class_name, class_bases, class_body):
+    class_dict = {}
+    for key, value in class_body.items():
+        class_dict[key] = value
+    return type(class_name, tuple(class_bases), class_dict)
+
+
+def evaluate_function_def(func_def, state, tools):
+    tools[func_def.name] = create_function(func_def, state, tools)
+    return tools[func_def.name]
+
+
+def evaluate_class_def(class_def, state, tools):
+    class_name = class_def.name
+    bases = [evaluate_ast(base, state, tools) for base in class_def.bases]
+    class_dict = {}
+
+    for stmt in class_def.body:
+        if isinstance(stmt, ast.FunctionDef):
+            class_dict[stmt.name] = evaluate_function_def(stmt, state, tools)
+        elif isinstance(stmt, ast.Assign):
+            for target in stmt.targets:
+                class_dict[target.id] = evaluate_ast(stmt.value, state, tools)
+        else:
+            raise InterpretorError(f"Unsupported statement in class body: {stmt.__class__.__name__}")
+
+    new_class = type(class_name, tuple(bases), class_dict)
+    state[class_name] = new_class
+    return new_class
 
 
 def evaluate_augassign(expression: ast.AugAssign, state: Dict[str, Any], tools: Dict[str, Callable]):
@@ -176,11 +236,20 @@ def evaluate_assign(assign, state, tools):
     var_names = assign.targets
     result = evaluate_ast(assign.value, state, tools)
     if len(var_names) == 1:
-        if isinstance(var_names[0], ast.Tuple):
-            for i, elem in enumerate(var_names[0].elts):
+        target = var_names[0]
+        if isinstance(target, ast.Tuple):
+            for i, elem in enumerate(target.elts):
                 state[elem.id] = result[i]
+        elif isinstance(target, ast.Attribute):
+            obj = evaluate_ast(target.value, state, tools)
+            setattr(obj, target.attr, result)
+        elif isinstance(target, ast.Subscript):
+            obj = evaluate_ast(target.value, state, tools)
+            key = evaluate_ast(target.slice, state, tools)
+            obj[key] = result
         else:
-            state[var_names[0].id] = result
+            state[target.id] = result
+
     else:
         if len(result) != len(var_names):
             raise InterpretorError(f"Expected {len(var_names)} values but got {len(result)}.")
@@ -190,41 +259,64 @@ def evaluate_assign(assign, state, tools):
 
 
 def evaluate_call(call, state, tools):
+    if not (isinstance(call.func, ast.Attribute) or isinstance(call.func, ast.Name)):
+        raise InterpretorError(
+            f"It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func})."
+        )
     if isinstance(call.func, ast.Attribute):
         obj = evaluate_ast(call.func.value, state, tools)
         func_name = call.func.attr
         if not hasattr(obj, func_name):
             raise InterpretorError(f"Object {obj} has no attribute {func_name}")
         func = getattr(obj, func_name)
-        args = [evaluate_ast(arg, state, tools) for arg in call.args]
-        kwargs = {keyword.arg: evaluate_ast(keyword.value, state, tools) for keyword in call.keywords}
-        return func(*args, **kwargs)
-
     elif isinstance(call.func, ast.Name):
         func_name = call.func.id
-
         if func_name in state:
             func = state[func_name]
         elif func_name in tools:
             func = tools[func_name]
+        elif func_name in ERRORS:
+            func = ERRORS[func_name]
         else:
             raise InterpretorError(
                 f"It is not permitted to evaluate other functions than the provided tools or imported functions (tried to execute {call.func.id})."
             )
-        # Todo deal with args
-        args = [evaluate_ast(arg, state, tools) for arg in call.args]
-        kwargs = {keyword.arg: evaluate_ast(keyword.value, state, tools) for keyword in call.keywords}
-        output = func(*args, **kwargs)
 
-        # store logs of print statements
-        if func_name == "print":
-            state["print_outputs"] += output + "\n"
+    args = [evaluate_ast(arg, state, tools) for arg in call.args]
+    kwargs = {keyword.arg: evaluate_ast(keyword.value, state, tools) for keyword in call.keywords}
 
-        return output
+    if isinstance(func, type) and len(func.__module__.split(".")) > 1:  # Check for user-defined classes
+        # Instantiate the class using its constructor
+        obj = func.__new__(func)  # Create a new instance of the class
+        if hasattr(obj, "__init__"):  # Check if the class has an __init__ method
+            obj.__init__(*args, **kwargs)  # Call the __init__ method correctly
+        return obj
     else:
-        raise InterpretorError(
-            f"It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func})."
-        )
+        if func_name == "super":
+            if not args:
+                if "__class__" in state and "self" in state:
+                    return super(state["__class__"], state["self"])
+                else:
+                    raise InterpretorError("super() needs at least one argument")
+            cls = args[0]
+            if not isinstance(cls, type):
+                raise InterpretorError("super() argument 1 must be type")
+            if len(args) == 1:
+                return super(cls)
+            elif len(args) == 2:
+                instance = args[1]
+                return super(cls, instance)
+            else:
+                raise InterpretorError("super() takes at most 2 arguments")
+
+        else:
+            if func_name == "print":
+                output = " ".join(map(str, args))
+                state["print_outputs"] += output + "\n"
+                return output
+            else:  # Assume it's a callable object
+                output = func(*args, **kwargs)
+                return output
 
 
 def evaluate_subscript(subscript, state, tools):
@@ -248,6 +340,10 @@ def evaluate_subscript(subscript, state, tools):
 def evaluate_name(name, state, tools):
     if name.id in state:
         return state[name.id]
+    elif name.id in tools:
+        return tools[name.id]
+    elif name.id in ERRORS:
+        return ERRORS[name.id]
     close_matches = difflib.get_close_matches(name.id, list(state.keys()))
     if len(close_matches) > 0:
         return state[close_matches[0]]
@@ -307,7 +403,11 @@ def evaluate_for(for_loop, state, tools):
     result = None
     iterator = evaluate_ast(for_loop.iter, state, tools)
     for counter in iterator:
-        state[for_loop.target.id] = counter
+        if isinstance(for_loop.target, ast.Tuple):
+            for i, elem in enumerate(for_loop.target.elts):
+                state[elem.id] = counter[i]
+        else:
+            state[for_loop.target.id] = counter
         for node in for_loop.body:
             try:
                 line_result = evaluate_ast(node, state, tools)
@@ -337,7 +437,56 @@ def evaluate_listcomp(listcomp, state, tools):
     return result
 
 
-def evaluate_ast(expression: ast.AST, state: Dict[str, Any], tools: Dict[str, Callable]):
+def evaluate_try(try_node, state, tools):
+    try:
+        for stmt in try_node.body:
+            evaluate_ast(stmt, state, tools)
+    except Exception as e:
+        matched = False
+        for handler in try_node.handlers:
+            if handler.type is None or isinstance(e, evaluate_ast(handler.type, state, tools)):
+                matched = True
+                if handler.name:
+                    state[handler.name] = e
+                for stmt in handler.body:
+                    evaluate_ast(stmt, state, tools)
+                break
+        if not matched:
+            raise e
+    else:
+        if try_node.orelse:
+            for stmt in try_node.orelse:
+                evaluate_ast(stmt, state, tools)
+    finally:
+        if try_node.finalbody:
+            for stmt in try_node.finalbody:
+                evaluate_ast(stmt, state, tools)
+
+
+def evaluate_raise(raise_node, state, tools):
+    if raise_node.exc is not None:
+        exc = evaluate_ast(raise_node.exc, state, tools)
+    else:
+        exc = None
+    if raise_node.cause is not None:
+        cause = evaluate_ast(raise_node.cause, state, tools)
+    else:
+        cause = None
+    if exc is not None:
+        if cause is not None:
+            raise exc from cause
+        else:
+            raise exc
+    else:
+        raise InterpretorError("Re-raise is not supported without an active exception")
+
+
+def evaluate_ast(
+    expression: ast.AST,
+    state: Dict[str, Any],
+    tools: Dict[str, Callable],
+    authorized_imports: List[str] = LIST_SAFE_MODULES,
+):
     """
     Evaluate an abstract syntax tree using the content of the variables stored in a state and only evaluating a given
     set of functions.
@@ -353,6 +502,9 @@ def evaluate_ast(expression: ast.AST, state: Dict[str, Any], tools: Dict[str, Ca
         tools (`Dict[str, Callable]`):
             The functions that may be called during the evaluation. Any call to another function will fail with an
             `InterpretorError`.
+        authorized_imports (`List[str]`):
+            The list of modules that can be imported by the code. By default, only a few safe modules are allowed.
+            Add more at your own risk!
     """
     if isinstance(expression, ast.Assign):
         # Assignement -> we evaluate the assignement which should update the state
@@ -459,7 +611,7 @@ def evaluate_ast(expression: ast.AST, state: Dict[str, Any], tools: Dict[str, Ca
         return result
     elif isinstance(expression, ast.Import):
         for alias in expression.names:
-            if alias.name in LIST_SAFE_MODULES:
+            if alias.name in authorized_imports:
                 module = __import__(alias.name)
                 state[alias.asname or alias.name] = module
             else:
@@ -468,19 +620,27 @@ def evaluate_ast(expression: ast.AST, state: Dict[str, Any], tools: Dict[str, Ca
     elif isinstance(expression, ast.While):
         return evaluate_while(expression, state, tools)
     elif isinstance(expression, ast.ImportFrom):
-        if expression.module in LIST_SAFE_MODULES:
+        if expression.module in authorized_imports:
             module = __import__(expression.module)
             for alias in expression.names:
                 state[alias.asname or alias.name] = getattr(module, alias.name)
         else:
             raise InterpretorError(f"Import from {expression.module} is not allowed.")
         return None
+    elif isinstance(expression, ast.ClassDef):
+        return evaluate_class_def(expression, state, tools)
+    elif isinstance(expression, ast.Try):
+        return evaluate_try(expression, state, tools)
+    elif isinstance(expression, ast.Raise):
+        return evaluate_raise(expression, state, tools)
     else:
         # For now we refuse anything else. Let's add things as we need them.
         raise InterpretorError(f"{expression.__class__.__name__} is not supported.")
 
 
-def evaluate_python_code(code: str, tools: Optional[Dict[str, Callable]] = {}, state=None):
+def evaluate_python_code(
+    code: str, tools: Optional[Dict[str, Callable]] = {}, state=None, authorized_imports: List[str] = LIST_SAFE_MODULES
+):
     """
     Evaluate a python expression using the content of the variables stored in a state and only evaluating a given set
     of functions.
@@ -506,9 +666,10 @@ def evaluate_python_code(code: str, tools: Optional[Dict[str, Callable]] = {}, s
         state = {}
     result = None
     state["print_outputs"] = ""
+
     for idx, node in enumerate(expression.body):
         try:
-            line_result = evaluate_ast(node, state, tools)
+            line_result = evaluate_ast(node, state, tools, authorized_imports)
         except InterpretorError as e:
             msg = f"You tried to execute the following code:\n{code}\n"
             msg += f"You got these outputs:\n{state['print_outputs']}\n"
diff --git a/src/transformers/agents/tools.py b/src/transformers/agents/tools.py
index 4016a20f81..36e88d5a06 100644
--- a/src/transformers/agents/tools.py
+++ b/src/transformers/agents/tools.py
@@ -185,7 +185,7 @@ class Tool:
             "tool_class": full_name,
             "description": self.description,
             "name": self.name,
-            "inputs": str(self.inputs),
+            "inputs": self.inputs,
             "output_type": str(self.output_type),
         }
         with open(config_file, "w", encoding="utf-8") as f:
@@ -315,7 +315,7 @@ class Tool:
         if tool_class.output_type != custom_tool["output_type"]:
             tool_class.output_type = custom_tool["output_type"]
 
-        return tool_class(model_repo_id, token=token, **kwargs)
+        return tool_class(**kwargs)
 
     def push_to_hub(
         self,
diff --git a/src/transformers/optimization.py b/src/transformers/optimization.py
index 79a2c71c38..a462e3d824 100644
--- a/src/transformers/optimization.py
+++ b/src/transformers/optimization.py
@@ -540,6 +540,9 @@ def get_scheduler(
     if name == SchedulerType.INVERSE_SQRT:
         return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
 
+    if name == SchedulerType.WARMUP_STABLE_DECAY:
+        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, **scheduler_specific_kwargs)
+
     # All other schedulers require `num_training_steps`
     if num_training_steps is None:
         raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py
index f204d4e02c..ef893acb6c 100644
--- a/src/transformers/quantizers/quantizer_bnb_4bit.py
+++ b/src/transformers/quantizers/quantizer_bnb_4bit.py
@@ -63,7 +63,7 @@ class Bnb4BitHfQuantizer(HfQuantizer):
         if not (is_accelerate_available() and is_bitsandbytes_available()):
             raise ImportError(
                 "Using `bitsandbytes` 8-bit quantization requires Accelerate: `pip install accelerate` "
-                "and the latest version of bitsandbytes: `pip install -i https://pypi.org/simple/ bitsandbytes`"
+                "and the latest version of bitsandbytes: `pip install -U bitsandbytes`"
             )
 
         if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py
index 906457f310..0573ce7532 100644
--- a/src/transformers/quantizers/quantizer_bnb_8bit.py
+++ b/src/transformers/quantizers/quantizer_bnb_8bit.py
@@ -64,7 +64,7 @@ class Bnb8BitHfQuantizer(HfQuantizer):
         if not (is_accelerate_available() and is_bitsandbytes_available()):
             raise ImportError(
                 "Using `bitsandbytes` 8-bit quantization requires Accelerate: `pip install accelerate` "
-                "and the latest version of bitsandbytes: `pip install -i https://pypi.org/simple/ bitsandbytes`"
+                "and the latest version of bitsandbytes: `pip install -U bitsandbytes`"
             )
 
         if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 49e7803066..e1871eaf8c 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -2175,6 +2175,9 @@ class Trainer:
         grad_norm: Optional[float] = None
         self.control = self.callback_handler.on_train_begin(args, self.state, self.control)
 
+        if args.sanity_evaluation:
+            self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True)
+
         total_batched_samples = 0
         for epoch in range(epochs_trained, num_train_epochs):
             epoch_iterator = train_dataloader
@@ -2723,6 +2726,18 @@ class Trainer:
                 f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}."
             )
 
+    def _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler=False):
+        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
+        self._report_to_hp_search(trial, self.state.global_step, metrics)
+
+        # Run delayed LR scheduler now that metrics are populated
+        if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) and not skip_scheduler:
+            metric_to_check = self.args.metric_for_best_model
+            if not metric_to_check.startswith("eval_"):
+                metric_to_check = f"eval_{metric_to_check}"
+            self.lr_scheduler.step(metrics[metric_to_check])
+        return metrics
+
     def _maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval):
         if self.control.should_log and self.state.global_step > self._globalstep_last_logged:
             if is_torch_xla_available():
@@ -2749,15 +2764,7 @@ class Trainer:
 
         metrics = None
         if self.control.should_evaluate:
-            metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
-            self._report_to_hp_search(trial, self.state.global_step, metrics)
-
-            # Run delayed LR scheduler now that metrics are populated
-            if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
-                metric_to_check = self.args.metric_for_best_model
-                if not metric_to_check.startswith("eval_"):
-                    metric_to_check = f"eval_{metric_to_check}"
-                self.lr_scheduler.step(metrics[metric_to_check])
+            metrics = self._evaluate(trial, ignore_keys_for_eval)
 
         if self.control.should_save:
             self._save_checkpoint(model, trial, metrics=metrics)
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
index a97139a07b..b506e4b150 100644
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -771,6 +771,9 @@ class TrainingArguments:
             rather than saving all eval logits in memory. When set to `True`, you must pass a compute_metrics function
             that takes a boolean argument `compute_result`, which when passed `True`, will trigger the final global
             summary statistics from the batch-level summary statistics you've accumulated over the evaluation set.
+
+        sanity_evaluation(`bool`, *optional*, defaults to `False`):
+            Whether or not to perform a sanity check to ensure that the validation steps works correctly. It will be performed before the training.
     """
 
     framework = "pt"
@@ -1454,6 +1457,13 @@ class TrainingArguments:
         metadata={"help": "Break eval metrics calculation into batches to save memory."},
     )
 
+    sanity_evaluation: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to run through the entire `evaluation` step at the very beginning of training as a sanity check."
+        },
+    )
+
     def __post_init__(self):
         # Parse in args that could be `dict` sent in from the CLI as a string
         for field in _VALID_DICT_FIELDS:
diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py
index f9e503cf86..6236827de3 100755
--- a/src/transformers/utils/quantization_config.py
+++ b/src/transformers/utils/quantization_config.py
@@ -383,6 +383,10 @@ class BitsAndBytesConfig(QuantizationConfigMixin):
         if bnb_4bit_quant_storage is None:
             self.bnb_4bit_quant_storage = torch.uint8
         elif isinstance(bnb_4bit_quant_storage, str):
+            if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]:
+                raise ValueError(
+                    "`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') "
+                )
             self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage)
         elif isinstance(bnb_4bit_quant_storage, torch.dtype):
             self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
diff --git a/utils/patch_helper.py b/utils/patch_helper.py
new file mode 100644
index 0000000000..86b32c9304
--- /dev/null
+++ b/utils/patch_helper.py
@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This should help you prepare a patch, automatically extracting the commits to cherry-pick
+in chronological order to avoid merge conflicts. An equivalent way to do this is to use
+`git log --pretty=oneline HEAD...v4.41.0` and grep.
+
+Potential TODO: automatically cherry-picks them.
+
+Pass in a list of PR:
+`python utils/patch_helper.py --prs 31108 31054 31008 31010 31004`
+will produce the following:
+```bash
+Skipping invalid version tag: list
+Skipping invalid version tag: localattn1
+Git cherry-pick commands to run:
+git cherry-pick 03935d300d60110bb86edb49d2315089cfb19789 #2024-05-24 11:00:59+02:00
+git cherry-pick bdb9106f247fca48a71eb384be25dbbd29b065a8 #2024-05-24 19:02:55+02:00
+git cherry-pick 84c4b72ee99e8e65a8a5754a5f9d6265b45cf67e #2024-05-27 10:34:14+02:00
+git cherry-pick 936ab7bae5e040ec58994cb722dd587b9ab26581 #2024-05-28 11:56:05+02:00
+git cherry-pick 0bef4a273825d2cfc52ddfe62ba486ee61cc116f #2024-05-29 13:33:26+01:00
+```
+"""
+
+import argparse
+
+from git import GitCommandError, Repo
+from packaging import version
+
+
+def get_merge_commit(repo, pr_number, since_tag):
+    try:
+        # Use git log to find the merge commit for the PR within the given tag range
+        merge_commit = next(repo.iter_commits(f"v{since_tag}...HEAD", grep=f"#{pr_number}"))
+        return merge_commit
+    except StopIteration:
+        print(f"No merge commit found for PR #{pr_number} between tags {since_tag} and {main}")
+        return None
+    except GitCommandError as e:
+        print(f"Error finding merge commit for PR #{pr_number}: {str(e)}")
+        return None
+
+
+def main(pr_numbers):
+    repo = Repo(".")  # Initialize the Repo object for the current directory
+    merge_commits = []
+
+    tags = {}
+    for tag in repo.tags:
+        try:
+            # Parse and sort tags, skip invalid ones
+            tag_ver = version.parse(tag.name)
+            tags[tag_ver] = tag
+        except Exception:
+            print(f"Skipping invalid version tag: {tag.name}")
+
+    last_tag = sorted(tags)[-1]
+    major_minor = f"{last_tag.major}.{last_tag.minor}.0"
+    # Iterate through tag ranges to find the merge commits
+    for pr in pr_numbers:
+        commit = get_merge_commit(repo, pr, major_minor)
+        if commit:
+            merge_commits.append(commit)
+
+    # Sort commits by date
+    merge_commits.sort(key=lambda commit: commit.committed_datetime)
+
+    # Output the git cherry-pick commands
+    print("Git cherry-pick commands to run:")
+    for commit in merge_commits:
+        print(f"git cherry-pick {commit.hexsha} #{commit.committed_datetime}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Find and sort merge commits for specified PRs.")
+    parser.add_argument("--prs", nargs="+", required=True, type=int, help="PR numbers to find merge commits for")
+
+    args = parser.parse_args()
+    main(args.prs)


Test patch:
diff --git a/tests/agents/test_python_interpreter.py b/tests/agents/test_python_interpreter.py
index 7a3257494f..dbe6c90a9e 100644
--- a/tests/agents/test_python_interpreter.py
+++ b/tests/agents/test_python_interpreter.py
@@ -353,3 +353,131 @@ if char.isalpha():
         result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state=state)
         assert result == "Ok no one cares"
         assert state["print_outputs"] == "Hello world!\nOk no one cares\n"
+
+    def test_tuple_target_in_iterator(self):
+        code = "for a, b in [('Ralf Weikert', 'Austria'), ('Samuel Seungwon Lee', 'South Korea')]:res = a.split()[0]"
+        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
+        assert result == "Samuel"
+
+    def test_classes(self):
+        code = """
+class Animal:
+    species = "Generic Animal"
+
+    def __init__(self, name, age):
+        self.name = name
+        self.age = age
+
+    def sound(self):
+        return "The animal makes a sound."
+
+    def __str__(self):
+        return f"{self.name}, {self.age} years old"
+
+class Dog(Animal):
+    species = "Canine"
+
+    def __init__(self, name, age, breed):
+        super().__init__(name, age)
+        self.breed = breed
+
+    def sound(self):
+        return "The dog barks."
+
+    def __str__(self):
+        return f"{self.name}, {self.age} years old, {self.breed}"
+
+class Cat(Animal):
+    def sound(self):
+        return "The cat meows."
+
+    def __str__(self):
+        return f"{self.name}, {self.age} years old, {self.species}"
+
+
+# Testing multiple instances
+dog1 = Dog("Fido", 3, "Labrador")
+dog2 = Dog("Buddy", 5, "Golden Retriever")
+
+# Testing method with built-in function
+animals = [dog1, dog2, Cat("Whiskers", 2)]
+num_animals = len(animals)
+
+# Testing exceptions in methods
+class ExceptionTest:
+    def method_that_raises(self):
+        raise ValueError("An error occurred")
+
+try:
+    exc_test = ExceptionTest()
+    exc_test.method_that_raises()
+except ValueError as e:
+    exception_message = str(e)
+
+
+# Collecting results
+dog1_sound = dog1.sound()
+dog1_str = str(dog1)
+dog2_sound = dog2.sound()
+dog2_str = str(dog2)
+cat = Cat("Whiskers", 2)
+cat_sound = cat.sound()
+cat_str = str(cat)
+    """
+        state = {}
+        evaluate_python_code(code, {"print": print, "len": len, "super": super, "str": str, "sum": sum}, state=state)
+
+        # Assert results
+        assert state["dog1_sound"] == "The dog barks."
+        assert state["dog1_str"] == "Fido, 3 years old, Labrador"
+        assert state["dog2_sound"] == "The dog barks."
+        assert state["dog2_str"] == "Buddy, 5 years old, Golden Retriever"
+        assert state["cat_sound"] == "The cat meows."
+        assert state["cat_str"] == "Whiskers, 2 years old, Generic Animal"
+        assert state["num_animals"] == 3
+        assert state["exception_message"] == "An error occurred"
+
+    def test_variable_args(self):
+        code = """
+def var_args_method(self, *args, **kwargs):
+    return sum(args) + sum(kwargs.values())
+
+var_args_method(1, 2, 3, x=4, y=5)
+"""
+        state = {}
+        result = evaluate_python_code(code, {"sum": sum}, state=state)
+        assert result == 15
+
+    def test_exceptions(self):
+        code = """
+def method_that_raises(self):
+    raise ValueError("An error occurred")
+
+try:
+    method_that_raises()
+except ValueError as e:
+    exception_message = str(e)
+    """
+        state = {}
+        evaluate_python_code(code, {"print": print, "len": len, "super": super, "str": str, "sum": sum}, state=state)
+        assert state["exception_message"] == "An error occurred"
+
+    def test_subscript(self):
+        code = "vendor = {'revenue': 31000, 'rent': 50312}; vendor['ratio'] = round(vendor['revenue'] / vendor['rent'], 2)"
+
+        state = {}
+        evaluate_python_code(code, {"min": min, "print": print, "round": round}, state=state)
+        assert state["vendor"] == {"revenue": 31000, "rent": 50312, "ratio": 0.62}
+
+    def test_print(self):
+        code = "print(min([1, 2, 3]))"
+        state = {}
+        result = evaluate_python_code(code, {"min": min, "print": print}, state=state)
+        assert result == "1"
+        assert state["print_outputs"] == "1\n"
+
+    def test_types_as_objects(self):
+        code = "type_a = float(2); type_b = str; type_c = int"
+        state = {}
+        result = evaluate_python_code(code, {"float": float, "str": str, "int": int}, state=state)
+        assert result == int
diff --git a/tests/optimization/test_optimization.py b/tests/optimization/test_optimization.py
index 6d6707db5a..5240b03779 100644
--- a/tests/optimization/test_optimization.py
+++ b/tests/optimization/test_optimization.py
@@ -36,6 +36,7 @@ if is_torch_available():
         get_inverse_sqrt_schedule,
         get_linear_schedule_with_warmup,
         get_polynomial_decay_schedule_with_warmup,
+        get_scheduler,
         get_wsd_schedule,
     )
 
@@ -176,6 +177,27 @@ class ScheduleInitTest(unittest.TestCase):
             lrs_2 = unwrap_and_save_reload_schedule(scheduler, self.num_steps)
             self.assertListEqual(lrs_1, lrs_2, msg=f"failed for {scheduler_func} in save and reload")
 
+    def test_get_scheduler(self):
+        test_params = [
+            {
+                "name": "warmup_stable_decay",
+                "optimizer": self.optimizer,
+                "num_warmup_steps": 2,
+                "scheduler_specific_kwargs": {"num_stable_steps": 1, "num_decay_steps": 3},
+            },
+            {
+                "name": "warmup_stable_decay",
+                "optimizer": self.optimizer,
+                "num_warmup_steps": 2,
+                "num_training_steps": 10,
+                "scheduler_specific_kwargs": {"num_stable_steps": 1, "num_decay_steps": 3},
+            },
+            {"name": "cosine", "optimizer": self.optimizer, "num_warmup_steps": 2, "num_training_steps": 10},
+        ]
+
+        for param in test_params:
+            self.assertTrue(get_scheduler(**param), msg=f"failed for {param['name']} in get_scheduler")
+
 
 class LambdaScheduleWrapper:
     """See https://github.com/huggingface/transformers/issues/21689"""
diff --git a/tests/quantization/bnb/test_4bit.py b/tests/quantization/bnb/test_4bit.py
index 443b1020a3..ac17979d17 100644
--- a/tests/quantization/bnb/test_4bit.py
+++ b/tests/quantization/bnb/test_4bit.py
@@ -303,6 +303,13 @@ class Bnb4BitTest(Base4bitTest):
         model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-small", load_in_4bit=True, device_map="auto")
         self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.float32)
 
+    def test_bnb_4bit_wrong_config(self):
+        r"""
+        Test whether creating a bnb config with unsupported values leads to errors.
+        """
+        with self.assertRaises(ValueError):
+            _ = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_storage="add")
+
 
 @require_bitsandbytes
 @require_accelerate
diff --git a/tests/quantization/quanto_integration/test_quanto.py b/tests/quantization/quanto_integration/test_quanto.py
index b5020f8030..fef5c30018 100644
--- a/tests/quantization/quanto_integration/test_quanto.py
+++ b/tests/quantization/quanto_integration/test_quanto.py
@@ -448,7 +448,7 @@ class QuantoKVCacheQuantizationTest(unittest.TestCase):
     def test_quantized_cache(self):
         EXPECTED_TEXT_COMPLETION = [
             "Simply put, the theory of relativity states that 1) the speed of light is the same for all observers, and 2) the laws of physics are the same for all observers.\nThe first part of the theory of relativity",
-            "My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, my fries, my burgers, my hot dogs, my sandwiches, my salads, my chicken, my fish",
+            "My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, my fries, my chicken, my burgers, my hot dogs, my sandwiches, my salads, my p",
         ]
 
         prompts = [