From f314d18863ccb58ce37b02d2f922ac9ed5178f5f Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Wed, 21 May 2025 04:11:07 +0400
Subject: [PATCH] feat: add also base models

---
 setup_env.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/setup_env.py b/setup_env.py
index 2544c51..f15d65f 100644
--- a/setup_env.py
+++ b/setup_env.py
@@ -50,6 +50,12 @@ SUPPORTED_HF_MODELS = {
     "tiiuae/Falcon-E-1B-Instruct": {
         "model_name": "Falcon-E-1B-Instruct",
     },
+    "tiiuae/Falcon-E-3B-Base": {
+        "model_name": "Falcon-E-3B-Base",
+    },
+    "tiiuae/Falcon-E-1B-Base": {
+        "model_name": "Falcon-E-1B-Base",
+    },
 }
 
 SUPPORTED_QUANT_TYPES = {