Remove device from Caduceus for workflows

helicalAI · Dec 17, 2024 · 0b6eac2 · 0b6eac2
1 parent 3762baf
commit 0b6eac2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -119,11 +119,11 @@ jobs:
 
       - name: Execute Caduceus
         run: |
-          python examples/run_models/run_caduceus.py ++device="cuda"
+          python examples/run_models/run_caduceus.py
     
       - name: Fine-tune Caduceus
         run: |
-          python examples/fine_tune_models/fine_tune_caduceus.py ++device="cuda"
+          python examples/fine_tune_models/fine_tune_caduceus.py
 
       - name: Execute benchmarking
         run: |

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -137,11 +137,11 @@ jobs:
 
       - name: Execute Caduceus
         run: |
-          python examples/run_models/run_caduceus.py ++device="cuda"
+          python examples/run_models/run_caduceus.py
       
       - name: Fine-tune Caduceus
         run: |
-          python examples/fine_tune_models/fine_tune_caduceus.py ++device="cuda"
+          python examples/fine_tune_models/fine_tune_caduceus.py
 
       - name: Execute benchmarking
         run: |

diff --git a/helical/models/caduceus/model.py b/helical/models/caduceus/model.py
@@ -105,6 +105,12 @@ def process_data(self, sequences: List[str], return_tensors: str="pt", padding:
         ----------
         sequences : list[str]
             The input DNA sequences to be processed.
+        return_tensors : str, optional, default="pt"
+            The return type of the processed data.
+        padding : str, optional, default="max_length"
+            The padding strategy to be used.
+        truncation : bool, optional, default=True
+            Whether to truncate the sequences or not.
 
         Returns
         ----------