chore: list specific compute_type

winstxnhdw · Sep 17, 2023 · 028a591 · 028a591
1 parent d508b1b
commit 028a591
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/capgen/types/options.py b/capgen/types/options.py
@@ -1,4 +1,4 @@
-from typing import TypedDict
+from typing import Literal, TypedDict
 
 
 class TranscriberOptions(TypedDict):
@@ -14,5 +14,15 @@ class TranscriberOptions(TypedDict):
     num_workers (int) : the number of workers
     """
     model_size_or_path: str
-    compute_type: str
     num_workers: int
+    compute_type: Literal[
+        'auto',
+        'int8',
+        'int8_float32',
+        'int8_float16',
+        'int8_bfloat16',
+        'int16'
+        'float16',
+        'bfloat16',
+        'float32',
+    ]