make max train size default to none

rungalileo · Oct 19, 2023 · 1d37f19 · 1d37f19
1 parent 212554b
commit 1d37f19
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/dataquality/integrations/seq2seq/formatter.py b/dataquality/integrations/seq2seq/formatter.py
@@ -1,14 +1,14 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Dict, Type
+from typing import Dict, Optional, Type
 
 
 @dataclass
 class BaseFormatter(ABC):
     name: str
     input_col: str
     target_col: str
-    max_train_size: int
+    max_train_size: Optional[int] = None
 
     @abstractmethod
     def format_sample(self, sample: Dict[str, str]) -> Dict[str, str]:
@@ -21,7 +21,6 @@ class DefaultFormatter(BaseFormatter):
     name: str = "default"
     input_col: str = "text"
     target_col: str = "label"
-    max_train_size: int = 100
 
     def format_sample(self, sample: Dict[str, str]) -> Dict[str, str]:
         """Base formatter is identity function"""