abacusai 0.70.0__tar.gz → 0.71.0__tar.gz

{abacusai-0.70.0/abacusai.egg-info → abacusai-0.71.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.1
 Name: abacusai
-Version: 0.70.0
+Version: 0.71.0
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI

{abacusai-0.70.0 → abacusai-0.71.0}/abacusai/__init__.py

@@ -4,4 +4,4 @@ from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "0.70.0"
+__version__ = "0.71.0"

abacusai-0.71.0/abacusai/api_class/dataset.py

@@ -0,0 +1,10 @@
+import dataclasses
+
+from .abstract import ApiClass
+
+
+@dataclasses.dataclass
+class ParsingConfig(ApiClass):
+    escape: str = dataclasses.field(default='"')
+    csv_delimiter: str = dataclasses.field(default=None)
+    file_path_with_schema: str = dataclasses.field(default=None)

{abacusai-0.70.0 → abacusai-0.71.0}/abacusai/api_class/enums.py

@@ -302,3 +302,10 @@ class PythonFunctionArgumentType(ApiEnum):
     MONITOR_ID = 'MONITOR_ID'
     BATCH_PREDICTION_ID = 'BATCH_PREDICTION_ID'
     DEPLOYMENT_ID = 'DEPLOYMENT_ID'
+
+
+class VectorStoreTextEncoder(ApiEnum):
+    OPENAI = 'OPENAI'
+    E5 = 'E5'
+    E5_SMALL = 'E5_SMALL'
+    SENTENCE_BERT = 'SENTENCE_BERT'

{abacusai-0.70.0 → abacusai-0.71.0}/abacusai/api_class/model.py

@@ -12,7 +12,7 @@ class TrainingConfig(ApiClass):
     _support_kwargs: bool = dataclasses.field(default=True, repr=False, init=False)
 
     kwargs: dict = dataclasses.field(default_factory=dict)
-    problem_type: enums.ProblemType = dataclasses.field(default=None)
+    problem_type: enums.ProblemType = dataclasses.field(default=None, repr=False, init=False)
     algorithm: str = dataclasses.field(default=None)
 
 
@@ -21,7 +21,6 @@ class PersonalizationTrainingConfig(TrainingConfig):
     """
     Training config for the PERSONALIZATION problem type
     Args:
-        problem_type (ProblemType): PERSONALIZATION
         objective (PersonalizationObjective): Ranking scheme used to select final best model.
         sort_objective (PersonalizationObjective): Ranking scheme used to sort models on the metrics page.
         training_mode (PersonalizationTrainingMode): whether to train in production or experimental mode.
@@ -57,7 +56,6 @@ class PersonalizationTrainingConfig(TrainingConfig):
 
     """
     # top-level params
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.PERSONALIZATION, repr=False, init=False)
     objective: enums.PersonalizationObjective = dataclasses.field(default=None)
     sort_objective: enums.PersonalizationObjective = dataclasses.field(default=None)
     training_mode: enums.PersonalizationTrainingMode = dataclasses.field(default=None)
@@ -104,12 +102,15 @@ class PersonalizationTrainingConfig(TrainingConfig):
     max_user_history_len_percentile: int = dataclasses.field(default=None)
     downsample_item_popularity_percentile: float = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.PERSONALIZATION
 
+
+@dataclasses.dataclass
 class RegressionTrainingConfig(TrainingConfig):
     """
     Training config for the PREDICTIVE_MODELING problem type
     Args:
-        problem_type (ProblemType): PREDICTIVE_MODELING
         objective (RegressionObjective): Ranking scheme used to select final best model.
         sort_objective (RegressionObjective): Ranking scheme used to sort models on the metrics page.
         tree_hpo_mode: (RegressionTreeHPOMode): Turning off Rapid Experimentation will take longer to train.
@@ -125,7 +126,7 @@ class RegressionTrainingConfig(TrainingConfig):
         test_row_indicator (str): Column indicating which rows to use for training (TRAIN) and testing (TEST). Validation (VAL) can also be specified.
         rebalance_classes (bool): Class weights are computed as the inverse of the class frequency from the training dataset when this option is selected as "Yes". It is useful when the classes in the dataset are unbalanced.
                                   Re-balancing classes generally boosts recall at the cost of precision on rare classes.
-        rare_class_augmentation_threshold (float): Augments any rare class whose relative frequency with respect to the most frequent class is less than this threshold.
+        rare_class_augmentation_threshold (float): Augments any rare class whose relative frequency with respect to the most frequent class is less than this threshold. Default = 0.1
         augmentation_strategy (RegressionAugmentationStrategy): Strategy to deal with class imbalance and data augmentation.
         training_rows_downsample_ratio (float): Uses this ratio to train on a sample of the dataset provided.
         active_labels_column (str): Specify a column to use as the active columns in a multi label setting.
@@ -150,7 +151,6 @@ class RegressionTrainingConfig(TrainingConfig):
         custom_metrics (list[str]): Registered custom metrics available for selection.
 
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.PREDICTIVE_MODELING, repr=False, init=False)
     objective: enums.RegressionObjective = dataclasses.field(default=None)
     sort_objective: enums.RegressionObjective = dataclasses.field(default=None)
     tree_hpo_mode: enums.RegressionTreeHPOMode = dataclasses.field(default=enums.RegressionTreeHPOMode.THOROUGH)
@@ -205,13 +205,15 @@ class RegressionTrainingConfig(TrainingConfig):
     custom_loss_functions: List[str] = dataclasses.field(default=None)
     custom_metrics: List[str] = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.PREDICTIVE_MODELING
+
 
 @dataclasses.dataclass
 class ForecastingTrainingConfig(TrainingConfig):
     """
     Training config for the FORECASTING problem type
     Args:
-        problem_type (ProblemType): FORECASTING
         prediction_length (int): How many timesteps in the future to predict.
         objective (ForecastingObjective): Ranking scheme used to select final best model.
         sort_objective (ForecastingObjective): Ranking scheme used to sort models on the metrics page.
@@ -272,7 +274,6 @@ class ForecastingTrainingConfig(TrainingConfig):
         custom_loss_functions (list[str]): Registered custom losses available for selection.
         custom_metrics (list[str]): Registered custom metrics available for selection.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.FORECASTING, repr=False, init=False)
     prediction_length: int = dataclasses.field(default=None)
     objective: enums.ForecastingObjective = dataclasses.field(default=None)
     sort_objective: enums.ForecastingObjective = dataclasses.field(default=None)
@@ -339,13 +340,15 @@ class ForecastingTrainingConfig(TrainingConfig):
     custom_loss_functions: List[str] = dataclasses.field(default=None)
     custom_metrics: List[str] = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.FORECASTING
+
 
 @dataclasses.dataclass
 class NamedEntityExtractionTrainingConfig(TrainingConfig):
     """
     Training config for the NAMED_ENTITY_EXTRACTION problem type
     Args:
-        problem_type (ProblemType): NAMED_ENTITY_EXTRACTION
         objective (NERObjective): Ranking scheme used to select final best model.
         sort_objective (NERObjective): Ranking scheme used to sort models on the metrics page.
         ner_model_type (NERModelType): Type of NER model to use.
@@ -357,7 +360,6 @@ class NamedEntityExtractionTrainingConfig(TrainingConfig):
         document_format (NLPDocumentFormat): Format of the input documents.
         include_longformer (bool): Whether to include the longformer model.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.NAMED_ENTITY_EXTRACTION, repr=False, init=False)
     objective: enums.NERObjective = dataclasses.field(default=None)
     sort_objective: enums.NERObjective = dataclasses.field(default=None)
     ner_model_type: enums.NERModelType = dataclasses.field(default=None)
@@ -372,13 +374,15 @@ class NamedEntityExtractionTrainingConfig(TrainingConfig):
     document_format: enums.NLPDocumentFormat = dataclasses.field(default=None)
     include_longformer: bool = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.NAMED_ENTITY_EXTRACTION
+
 
 @dataclasses.dataclass
 class NaturalLanguageSearchTrainingConfig(TrainingConfig):
     """
     Training config for the NATURAL_LANGUAGE_SEARCH problem type
     Args:
-        problem_type (ProblemType): NATURAL_LANGUAGE_SEARCH
         abacus_internal_model (bool): Use a Abacus.AI LLM to answer questions about your data without using any external APIs
         num_completion_tokens (int): Default for maximum number of tokens for chat answers. Reducing this will get faster responses which are more succinct
         larger_embeddings (bool): Use a higher dimension embedding model.
@@ -386,7 +390,6 @@ class NaturalLanguageSearchTrainingConfig(TrainingConfig):
         chunk_overlap_fraction (float): Overlap in chunks while indexing the documents.
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.NATURAL_LANGUAGE_SEARCH, repr=False, init=False)
     abacus_internal_model: bool = dataclasses.field(default=None)
     num_completion_tokens: int = dataclasses.field(default=None)
     larger_embeddings: bool = dataclasses.field(default=None)
@@ -394,23 +397,29 @@ class NaturalLanguageSearchTrainingConfig(TrainingConfig):
     index_fraction: float = dataclasses.field(default=None)
     chunk_overlap_fraction: float = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.NATURAL_LANGUAGE_SEARCH
+
 
 @dataclasses.dataclass
 class ChatLLMTrainingConfig(TrainingConfig):
     """
     Training config for the CHAT_LLM problem type
     Args:
-        problem_type (ProblemType): CHAT_LLM
         llm_name (str):
-        num_completion_tokens (int): Default for maximum number of tokens for chat answers. Reducing this will get faster responses which are more succinct
-        search_chunk_size (int): Chunk size for indexing the documents.
-        chunk_overlap_fraction (float): Overlap in chunks while indexing the documents.
         vector_stores (List[str]): List of vector store names to use for the feature stores this model was trained with.
+        num_completion_tokens (int): Default for maximum number of tokens for chat answers. Reducing this will get faster responses which are more succinct
+        system_message (str): The generative LLM system message
+        temperature (float): The generative LLM temperature
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.CHAT_LLM, repr=False, init=False)
-    llm_name: str = dataclasses.field(default=None)
-    num_completion_tokens: int = dataclasses.field(default=None)
-    vector_stores: List[str] = dataclasses.field(default=None)
+    llm_name: str = None
+    vector_stores: List[str] = None
+    num_completion_tokens: int = None
+    system_message: str = None
+    temperature: float = None
+
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CHAT_LLM
 
 
 @dataclasses.dataclass
@@ -418,32 +427,31 @@ class SentenceBoundaryDetectionTrainingConfig(TrainingConfig):
     """
     Training config for the SENTENCE_BOUNDARY_DETECTION problem type
     Args:
-        problem_type (ProblemType): SENTENCE_BOUNDARY_DETECTION
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
         dropout_rate (float): Dropout rate for neural network.
         batch_size (BatchSize): Batch size for neural network.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.SENTENCE_BOUNDARY_DETECTION, repr=False, init=False)
     # Data Split Params
     test_split: int = dataclasses.field(default=None)
     # Neural Network
     dropout_rate: float = dataclasses.field(default=None)
     batch_size: enums.BatchSize = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.SENTENCE_BOUNDARY_DETECTION
+
 
 @dataclasses.dataclass
 class SentimentDetectionTrainingConfig(TrainingConfig):
     """
     Training config for the SENTIMENT_DETECTION problem type
     Args:
-        problem_type (ProblemType): SENTIMENT_DETECTION
         sentiment_type (SentimentType): Type of sentiment to detect.
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
         dropout_rate (float): Dropout rate for neural network.
         batch_size (BatchSize): Batch size for neural network.
         compute_metrics (bool): Whether to compute metrics.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.SENTIMENT_DETECTION, repr=False, init=False)
     sentiment_type: enums.SentimentType = dataclasses.field(default=None)
     # Data Split Params
     test_split: int = dataclasses.field(default=None)
@@ -453,19 +461,20 @@ class SentimentDetectionTrainingConfig(TrainingConfig):
     # Metrics
     compute_metrics: bool = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.SENTIMENT_DETECTION
+
 
 @dataclasses.dataclass
 class DocumentClassificationTrainingConfig(TrainingConfig):
     """
     Training config for the DOCUMENT_CLASSIFICATION problem type
     Args:
-        problem_type (ProblemType): DOCUMENT_CLASSIFICATION
         zero_shot_hypotheses (List[str]): Zero shot hypotheses. Example text: 'This text is about pricing'.
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
         dropout_rate (float): Dropout rate for neural network.
         batch_size (BatchSize): Batch size for neural network.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.DOCUMENT_CLASSIFICATION, repr=False, init=False)
     zero_shot_hypotheses: List[str] = dataclasses.field(default=None)
     # Data Split Params
     test_split: int = dataclasses.field(default=None)
@@ -473,75 +482,81 @@ class DocumentClassificationTrainingConfig(TrainingConfig):
     dropout_rate: float = dataclasses.field(default=None)
     batch_size: enums.BatchSize = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.DOCUMENT_CLASSIFICATION
+
 
 @dataclasses.dataclass
 class DocumentSummarizationTrainingConfig(TrainingConfig):
     """
     Training config for the DOCUMENT_SUMMARIZATION problem type
     Args:
-        problem_type (ProblemType): DOCUMENT_SUMMARIZATION
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
         dropout_rate (float): Dropout rate for neural network.
         batch_size (BatchSize): Batch size for neural network.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.DOCUMENT_SUMMARIZATION, repr=False, init=False)
     # Data Split Params
     test_split: int = dataclasses.field(default=None)
     # Neural Network
     dropout_rate: float = dataclasses.field(default=None)
     batch_size: enums.BatchSize = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.DOCUMENT_SUMMARIZATION
+
 
 @dataclasses.dataclass
 class DocumentVisualizationTrainingConfig(TrainingConfig):
     """
     Training config for the DOCUMENT_VISUALIZATION problem type
     Args:
-        problem_type (ProblemType): DOCUMENT_VISUALIZATION
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
         dropout_rate (float): Dropout rate for neural network.
         batch_size (BatchSize): Batch size for neural network.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.DOCUMENT_VISUALIZATION, repr=False, init=False)
     # Data Split Params
     test_split: int = dataclasses.field(default=None)
     # Neural Network
     dropout_rate: float = dataclasses.field(default=None)
     batch_size: enums.BatchSize = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.DOCUMENT_VISUALIZATION
+
 
 @dataclasses.dataclass
 class ClusteringTrainingConfig(TrainingConfig):
     """
     Training config for the CLUSTERING problem type
     Args:
-        problem_type (ProblemType): CLUSTERING
         num_clusters_selection (int): Number of clusters. If None, will be selected automatically.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.CLUSTERING, repr=False, init=False)
     num_clusters_selection: int = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CLUSTERING
+
 
 @dataclasses.dataclass
 class ClusteringTimeseriesTrainingConfig(TrainingConfig):
     """
     Training config for the CLUSTERING_TIMESERIES problem type
     Args:
-        problem_type (ProblemType): CLUSTERING_TIMESERIES
         num_clusters_selection (int): Number of clusters. If None, will be selected automatically.
         imputation (ClusteringImputationMethod): Imputation method for missing values.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.CLUSTERING_TIMESERIES, repr=False, init=False)
     num_clusters_selection: int = dataclasses.field(default=None)
     imputation: enums.ClusteringImputationMethod = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CLUSTERING_TIMESERIES
+
 
 @dataclasses.dataclass
 class CumulativeForecastingTrainingConfig(TrainingConfig):
     """
     Training config for the CUMULATIVE_FORECASTING problem type
     Args:
-        problem_type (ProblemType): CUMULATIVE_FORECASTING
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
         historical_frequency (str): Forecast frequency
         cumulative_prediction_lengths (List[int]): List of Cumulative Prediction Frequencies. Each prediction length must be between 1 and 365.
@@ -549,7 +564,6 @@ class CumulativeForecastingTrainingConfig(TrainingConfig):
         skip_target_transform (bool): Avoid doing numeric scaling transformations on the target.
         predict_residuals (bool): Predict residuals instead of totals at each prediction step.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.CUMULATIVE_FORECASTING, repr=False, init=False)
     test_split: int = dataclasses.field(default=None)
     historical_frequency: str = dataclasses.field(default=None)
     cumulative_prediction_lengths: List[int] = dataclasses.field(default=None)
@@ -557,12 +571,15 @@ class CumulativeForecastingTrainingConfig(TrainingConfig):
     skip_target_transform: bool = dataclasses.field(default=None)
     predict_residuals: bool = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CUMULATIVE_FORECASTING
 
+
+@dataclasses.dataclass
 class AnomalyDetectionTrainingConfig(TrainingConfig):
     """
     Training config for the ANOMALY_DETECTION problem type
     Args:
-        problem_type (ProblemType): ANOMALY_DETECTION
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 (i.e. 5%) to 20 (i.e. 20%) of your dataset to use as test data.
         value_high (bool): Detect unusually high values.
         mixture_of_gaussians (bool): Detect unusual combinations of values using mixture of Gaussians.
@@ -571,7 +588,6 @@ class AnomalyDetectionTrainingConfig(TrainingConfig):
         spike_down (bool): Detect outliers with a low value.
         trend_change (bool): Detect changes to the trend.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.ANOMALY_DETECTION, repr=False, init=False)
     test_split: int = dataclasses.field(default=None)
     value_high: bool = dataclasses.field(default=None)
     mixture_of_gaussians: bool = dataclasses.field(default=None)
@@ -580,54 +596,59 @@ class AnomalyDetectionTrainingConfig(TrainingConfig):
     spike_down: bool = dataclasses.field(default=None)
     trend_change: bool = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.ANOMALY_DETECTION
+
 
 @dataclasses.dataclass
 class ThemeAnalysisTrainingConfig(TrainingConfig):
     """
     Training config for the THEME ANALYSIS problem type
-    Args:
-        problem_type (ProblemType): THEME_ANALYSIS
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.THEME_ANALYSIS, repr=False, init=False)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.THEME_ANALYSIS
 
+
+@dataclasses.dataclass
 class AIAgentTrainingConfig(TrainingConfig):
     """
     Training config for the AI_AGENT problem type
     Args:
-        problem_type (ProblemType): AI_AGENT
         description (str): Description of the agent function.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.AI_AGENT, repr=False, init=False)
     description: str = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.AI_AGENT
+
 
 @dataclasses.dataclass
 class CustomTrainedModelTrainingConfig(TrainingConfig):
     """
     Training config for the CUSTOM_TRAINED_MODEL problem type
     Args:
-        problem_type (ProblemType): CUSTOM_TRAINED_MODEL
         max_catalog_size (int): Maximum expected catalog size.
         max_dimension (int): Maximum expected dimension of the catalog.
         index_output_path (str): Fully qualified cloud location (GCS, S3, etc) to export snapshots of the embedding to.
         docker_image_uri (str): Docker image URI.
         service_port (int): Service port.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.CUSTOM_TRAINED_MODEL, repr=False, init=False)
     max_catalog_size: int = dataclasses.field(default=None)
     max_dimension: int = dataclasses.field(default=None)
     index_output_path: str = dataclasses.field(default=None)
     docker_image_uri: str = dataclasses.field(default=None)
     service_port: int = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CUSTOM_TRAINED_MODEL
+
 
 @dataclasses.dataclass
 class CustomAlgorithmTrainingConfig(TrainingConfig):
     """
     Training config for the CUSTOM_ALGORITHM problem type
     Args:
-        problem_type (ProblemType): CUSTOM_ALGORITHM
         train_function_name (str): The name of the train function.
         predict_many_function_name (str): The name of the predict many function.
         training_input_tables (List[str]): List of tables to use for training.
@@ -636,7 +657,6 @@ class CustomAlgorithmTrainingConfig(TrainingConfig):
         predict_module_name (str): The name of the predict module - only relevant if model is being uploaded from a zip file or github repositoty.
         test_split (int): Percent of dataset to use for test data. We support using a range between 6% to 20% of your dataset to use as test data.
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.CUSTOM_ALGORITHM, repr=False, init=False)
     train_function_name: str = dataclasses.field(default=None)
     predict_many_function_name: str = dataclasses.field(default=None)
     training_input_tables: List[str] = dataclasses.field(default=None)
@@ -645,15 +665,18 @@ class CustomAlgorithmTrainingConfig(TrainingConfig):
     predict_module_name: str = dataclasses.field(default=None)
     test_split: int = dataclasses.field(default=None)
 
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CUSTOM_ALGORITHM
+
 
 @dataclasses.dataclass
 class OptimizationTrainingConfig(TrainingConfig):
     """
     Training config for the OPTIMIZATION problem type
-    Args:
-        problem_type (ProblemType): OPTIMIZATION
     """
-    problem_type: enums.ProblemType = dataclasses.field(default=enums.ProblemType.OPTIMIZATION, repr=False, init=False)
+
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.OPTIMIZATION
 
 
 @dataclasses.dataclass

{abacusai-0.70.0 → abacusai-0.71.0}/abacusai/api_class/vector_store.py

@@ -1,6 +1,7 @@
 import dataclasses
 
 from .abstract import ApiClass
+from .enums import VectorStoreTextEncoder
 
 
 @dataclasses.dataclass
@@ -14,3 +15,11 @@ class VectorStoreConfig(ApiClass):
     """
     chunk_size: int = dataclasses.field(default=512)
     chunk_overlap_fraction: float = dataclasses.field(default=0.1)
+    text_encoder: VectorStoreTextEncoder = dataclasses.field(default=VectorStoreTextEncoder.OPENAI)
+
+
+@dataclasses.dataclass
+class DocumentRetrieverConfig(VectorStoreConfig):
+    """
+    Configs for document retriever.
+    """