[model upload] Support for organizations

src/transformers/hf_api.py

@@ -17,7 +17,7 @@
 import io
 import os
 from os.path import expanduser
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import requests
 from tqdm import tqdm
@@ -109,7 +109,7 @@ class HfApi:
         d = r.json()
         return d["token"]
 
-    def whoami(self, token: str) -> str:
+    def whoami(self, token: str) -> Tuple[str, List[str]]:
         """
         Call HF API to know "whoami"
         """
@@ -117,7 +117,7 @@ class HfApi:
         r = requests.get(path, headers={"authorization": "Bearer {}".format(token)})
         r.raise_for_status()
         d = r.json()
-        return d["user"]
+        return d["user"], d["orgs"]
 
     def logout(self, token: str) -> None:
         """
@@ -127,24 +127,28 @@ class HfApi:
         r = requests.post(path, headers={"authorization": "Bearer {}".format(token)})
         r.raise_for_status()
 
-    def presign(self, token: str, filename: str) -> PresignedUrl:
+    def presign(self, token: str, filename: str, organization: Optional[str] = None) -> PresignedUrl:
         """
         Call HF API to get a presigned url to upload `filename` to S3.
         """
         path = "{}/api/presign".format(self.endpoint)
-        r = requests.post(path, headers={"authorization": "Bearer {}".format(token)}, json={"filename": filename})
+        r = requests.post(
+            path,
+            headers={"authorization": "Bearer {}".format(token)},
+            json={"filename": filename, "organization": organization},
+        )
         r.raise_for_status()
         d = r.json()
         return PresignedUrl(**d)
 
-    def presign_and_upload(self, token: str, filename: str, filepath: str) -> str:
+    def presign_and_upload(self, token: str, filename: str, filepath: str, organization: Optional[str] = None) -> str:
         """
         Get a presigned url, then upload file to S3.
 
         Outputs:
             url: Read-only url for the stored file on S3.
         """
-        urls = self.presign(token, filename=filename)
+        urls = self.presign(token, filename=filename, organization=organization)
         # streaming upload:
         # https://2.python-requests.org/en/master/user/advanced/#streaming-uploads
         #
@@ -159,22 +163,27 @@ class HfApi:
             pf.close()
         return urls.access
 
-    def list_objs(self, token: str) -> List[S3Obj]:
+    def list_objs(self, token: str, organization: Optional[str] = None) -> List[S3Obj]:
         """
-        Call HF API to list all stored files for user.
+        Call HF API to list all stored files for user (or one of their organizations).
         """
         path = "{}/api/listObjs".format(self.endpoint)
-        r = requests.get(path, headers={"authorization": "Bearer {}".format(token)})
+        params = {"organization": organization} if organization is not None else None
+        r = requests.get(path, params=params, headers={"authorization": "Bearer {}".format(token)})
         r.raise_for_status()
         d = r.json()
         return [S3Obj(**x) for x in d]
 
-    def delete_obj(self, token: str, filename: str):
+    def delete_obj(self, token: str, filename: str, organization: Optional[str] = None):
         """
         Call HF API to delete a file stored by user
         """
         path = "{}/api/deleteObj".format(self.endpoint)
-        r = requests.delete(path, headers={"authorization": "Bearer {}".format(token)}, json={"filename": filename})
+        r = requests.delete(
+            path,
+            headers={"authorization": "Bearer {}".format(token)},
+            json={"filename": filename, "organization": organization},
+        )
         r.raise_for_status()
 
     def model_list(self) -> List[ModelInfo]: