Upgrade to SynapseAI 1.18 (#227)

Signed-off-by: yuanwu <[email protected]>
Co-authored-by: Thanaji Rao Thakkalapelli <[email protected]>

Dockerfile

@@ -32,7 +32,7 @@ COPY launcher launcher
 RUN cargo build --release
 
 # Text Generation Inference base image
-FROM vault.habana.ai/gaudi-docker/1.17.0/ubuntu22.04/habanalabs/pytorch-installer-2.3.1:latest as base
+FROM vault.habana.ai/gaudi-docker/1.18.0/ubuntu22.04/habanalabs/pytorch-installer-2.4.0:latest as base
 
 # Text Generation Inference base env
 ENV HUGGINGFACE_HUB_CACHE=/data \
@@ -61,7 +61,7 @@ RUN cd server && \
     make gen-server && \
     pip install -r requirements.txt && \
     bash ./dill-0.3.8-patch.sh && \
-    pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.17.0 && \
+    pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.18.0 && \
     BUILD_CUDA_EXT=0 pip install git+https://github.com/AutoGPTQ/AutoGPTQ.git@097dd04e --no-build-isolation && \
     pip install . --no-cache-dir
 

examples/run_generation.py

@@ -31,13 +31,18 @@ def get_args():
     parser.add_argument(
         "--max_concurrent_requests", type=int, default=256, help="Max number of concurrent requests"
     )
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for datasets"
+    )
+
     return parser.parse_args()
 
 
 def read_dataset(
     max_input_length: int,
     total_sample_count: int,
-    model_id: str
+    model_id: str,
+    seed: int,
 ) -> List[str]:
     """
     Loads public dataset from HF: https://huggingface.co/datasets/DIBT/10k_prompts_ranked
@@ -51,7 +56,7 @@ def read_dataset(
     )
     if len(dataset) > total_sample_count:
         dataset = dataset.select(range(total_sample_count))
-    dataset = dataset.shuffle()
+    dataset = dataset.shuffle(seed=seed)
     return [sample["prompt"] for sample in dataset]
 
 
@@ -71,7 +76,7 @@ def is_tgi_available(
 def main():
     args = get_args()
     dataset = read_dataset(
-        args.max_input_length, args.total_sample_count, args.model_id
+        args.max_input_length, args.total_sample_count, args.model_id, args.seed
     )
 
     if not is_tgi_available(args.server_address):