fixing python library requirements

kubeflow · Jan 4, 2024 · 95173aa · 95173aa
1 parent bf04ddc
commit 95173aa
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/sdk/python/kubeflow/storage_init_container/s3.py b/sdk/python/kubeflow/storage_init_container/s3.py
@@ -5,6 +5,7 @@
 from urllib.parse import urlparse
 import os
 
+
 @dataclass
 class S3DatasetParams:
     endpoint_url: str
@@ -50,6 +51,8 @@ def download_dataset(self):
 
         # Download the file
         s3_client.download_file(
-            self.config.bucket_name, self.config.file_key, os.path.join(self.config.download_dir,self.config.file_key)
+            self.config.bucket_name,
+            self.config.file_key,
+            os.path.join(self.config.download_dir, self.config.file_key),
         )
         print(f"File downloaded to: {self.config.download_dir}")
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
@@ -31,6 +31,11 @@
     "retrying>=1.3.3",
     "boto3>=1.33.9",
     "transformers>=4.35.2",
+    "torch>=2.1.1",
+    "torchvision>=0.16.1",
+    "torchaudio>=2.1.1",
+    "einops>=0.7.0",
+    "transformers_stream_generator>=0.0.4",
     "peft>=0.7.0",
 ]