orca-sdk 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

orca_sdk/datasource.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
 import logging
-import os
 import tempfile
 import zipfile
 from datetime import datetime
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import Literal, Union, cast
+from typing import Any, Literal, Union, cast
 
 import pandas as pd
 import pyarrow as pa
@@ -488,6 +487,50 @@ class Datasource:
     def __len__(self) -> int:
         return self.length
 
+    def query(
+        self,
+        offset: int = 0,
+        limit: int = 100,
+        shuffle: bool = False,
+        shuffle_seed: int | None = None,
+        filters: list[tuple[str, Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "like"], Any]] = [],
+    ) -> list[dict[str, Any]]:
+        """
+        Query the datasource for rows with pagination and filtering support.
+
+        Params:
+            offset: Number of rows to skip
+            limit: Maximum number of rows to return
+            shuffle: Whether to shuffle the dataset before pagination
+            shuffle_seed: Seed for shuffling (for reproducible results)
+            filters: List of filter tuples. Each tuple contains:
+                - field (str): Column name to filter on
+                - op (str): Operator ("==", "!=", ">", ">=", "<", "<=", "in", "not in", "like")
+                - value: Value to compare against
+
+        Returns:
+            List of rows from the datasource
+
+        Examples:
+            >>> datasource.query(filters=[("age", ">", 25)])
+            >>> datasource.query(filters=[("city", "in", ["NYC", "LA"])])
+            >>> datasource.query(filters=[("name", "like", "John")])
+        """
+
+        client = OrcaClient._resolve_client()
+        response = client.POST(
+            "/datasource/{name_or_id}/rows",
+            params={"name_or_id": self.id},
+            json={
+                "limit": limit,
+                "offset": offset,
+                "shuffle": shuffle,
+                "shuffle_seed": shuffle_seed,
+                "filters": [{"field": field, "op": op, "value": value} for field, op, value in filters],
+            },
+        )
+        return response
+
     def download(
         self, output_dir: str | PathLike, file_type: Literal["hf_dataset", "json", "csv"] = "hf_dataset"
     ) -> None:

orca_sdk/datasource_test.py

@@ -301,6 +301,126 @@ def test_from_disk_already_exists():
             os.unlink(f.name)
 
 
+def test_query_datasource_rows():
+    """Test querying rows from a datasource with pagination and shuffle."""
+    # Create a new dataset with 5 entries for testing
+    test_data = [{"id": i, "name": f"item_{i}"} for i in range(5)]
+    datasource = Datasource.from_list(name="test_query_datasource", data=test_data)
+
+    # Test basic query
+    rows = datasource.query(limit=3)
+    assert len(rows) == 3
+    assert all(isinstance(row, dict) for row in rows)
+
+    # Test offset
+    offset_rows = datasource.query(offset=2, limit=2)
+    assert len(offset_rows) == 2
+    assert offset_rows[0]["id"] == 2
+
+    # Test shuffle
+    shuffled_rows = datasource.query(limit=5, shuffle=True)
+    assert len(shuffled_rows) == 5
+    assert not all(row["id"] == i for i, row in enumerate(shuffled_rows))
+
+    # Test shuffle with seed
+    assert datasource.query(limit=5, shuffle=True, shuffle_seed=42) == datasource.query(
+        limit=5, shuffle=True, shuffle_seed=42
+    )
+
+
+def test_query_datasource_with_filters():
+    """Test querying datasource rows with various filter operators."""
+    # Create a datasource with test data
+    test_data = [
+        {"name": "Alice", "age": 25, "city": "New York", "score": 85.5},
+        {"name": "Bob", "age": 30, "city": "San Francisco", "score": 90.0},
+        {"name": "Charlie", "age": 35, "city": "Chicago", "score": 75.5},
+        {"name": "Diana", "age": 28, "city": "Boston", "score": 88.0},
+        {"name": "Eve", "age": 32, "city": "New York", "score": 92.0},
+    ]
+    datasource = Datasource.from_list(name=f"test_filter_datasource_{uuid4()}", data=test_data)
+
+    # Test == operator
+    rows = datasource.query(filters=[("city", "==", "New York")])
+    assert len(rows) == 2
+    assert all(row["city"] == "New York" for row in rows)
+
+    # Test > operator
+    rows = datasource.query(filters=[("age", ">", 30)])
+    assert len(rows) == 2
+    assert all(row["age"] > 30 for row in rows)
+
+    # Test >= operator
+    rows = datasource.query(filters=[("score", ">=", 88.0)])
+    assert len(rows) == 3
+    assert all(row["score"] >= 88.0 for row in rows)
+
+    # Test < operator
+    rows = datasource.query(filters=[("age", "<", 30)])
+    assert len(rows) == 2
+    assert all(row["age"] < 30 for row in rows)
+
+    # Test in operator
+    rows = datasource.query(filters=[("city", "in", ["New York", "Boston"])])
+    assert len(rows) == 3
+    assert all(row["city"] in ["New York", "Boston"] for row in rows)
+
+    # Test not in operator
+    rows = datasource.query(filters=[("city", "not in", ["New York", "Boston"])])
+    assert len(rows) == 2
+    assert all(row["city"] not in ["New York", "Boston"] for row in rows)
+
+    # Test like operator
+    rows = datasource.query(filters=[("name", "like", "li")])
+    assert len(rows) == 2
+    assert all("li" in row["name"].lower() for row in rows)
+
+    # Test multiple filters (AND logic)
+    rows = datasource.query(filters=[("city", "==", "New York"), ("age", ">", 26)])
+    assert len(rows) == 1
+    assert rows[0]["name"] == "Eve"
+
+    # Test filter with pagination
+    rows = datasource.query(filters=[("age", ">=", 28)], limit=2, offset=1)
+    assert len(rows) == 2
+
+
+def test_query_datasource_with_none_filters():
+    """Test filtering for None values."""
+    test_data = [
+        {"name": "Alice", "age": 25, "label": "A"},
+        {"name": "Bob", "age": 30, "label": None},
+        {"name": "Charlie", "age": 35, "label": "C"},
+        {"name": "Diana", "age": None, "label": "D"},
+        {"name": "Eve", "age": 32, "label": None},
+    ]
+    datasource = Datasource.from_list(name=f"test_none_filter_{uuid4()}", data=test_data)
+
+    # Test == None
+    rows = datasource.query(filters=[("label", "==", None)])
+    assert len(rows) == 2
+    assert all(row["label"] is None for row in rows)
+
+    # Test != None
+    rows = datasource.query(filters=[("label", "!=", None)])
+    assert len(rows) == 3
+    assert all(row["label"] is not None for row in rows)
+
+    # Test that None values are excluded from comparison operators
+    rows = datasource.query(filters=[("age", ">", 25)])
+    assert len(rows) == 3
+    assert all(row["age"] is not None and row["age"] > 25 for row in rows)
+
+
+def test_query_datasource_filter_invalid_column():
+    """Test that querying with an invalid column raises an error."""
+    test_data = [{"name": "Alice", "age": 25}]
+    datasource = Datasource.from_list(name=f"test_invalid_filter_{uuid4()}", data=test_data)
+
+    with pytest.raises(ValueError):
+        datasource.query(filters=[("invalid_column", "==", "test")])
+
+
 def test_to_list(hf_dataset, datasource):
     assert datasource.to_list() == hf_dataset.to_list()
 

orca_sdk/embedding_model.py

@@ -20,7 +20,7 @@ from .datasource import Datasource
 from .job import Job, Status
 
 if TYPE_CHECKING:
-    from .memoryset import LabeledMemoryset
+    from .memoryset import LabeledMemoryset, ScoredMemoryset
 
 
 class EmbeddingModelBase(ABC):
@@ -110,7 +110,7 @@ class EmbeddingModelBase(ABC):
         label_column: str,
         score_column: None = None,
         eval_datasource: Datasource | None = None,
-        subsample: int | None = None,
+        subsample: int | float | None = None,
         neighbor_count: int = 5,
         batch_size: int = 32,
         weigh_memories: bool = True,
@@ -127,7 +127,7 @@ class EmbeddingModelBase(ABC):
         label_column: str,
         score_column: None = None,
         eval_datasource: Datasource | None = None,
-        subsample: int | None = None,
+        subsample: int | float | None = None,
         neighbor_count: int = 5,
         batch_size: int = 32,
         weigh_memories: bool = True,
@@ -144,7 +144,7 @@ class EmbeddingModelBase(ABC):
         label_column: None = None,
         score_column: str,
         eval_datasource: Datasource | None = None,
-        subsample: int | None = None,
+        subsample: int | float | None = None,
         neighbor_count: int = 5,
         batch_size: int = 32,
         weigh_memories: bool = True,
@@ -161,7 +161,7 @@ class EmbeddingModelBase(ABC):
         label_column: None = None,
         score_column: str,
         eval_datasource: Datasource | None = None,
-        subsample: int | None = None,
+        subsample: int | float | None = None,
         neighbor_count: int = 5,
         batch_size: int = 32,
         weigh_memories: bool = True,
@@ -177,7 +177,7 @@ class EmbeddingModelBase(ABC):
         label_column: str | None = None,
         score_column: str | None = None,
         eval_datasource: Datasource | None = None,
-        subsample: int | None = None,
+        subsample: int | float | None = None,
         neighbor_count: int = 5,
         batch_size: int = 32,
         weigh_memories: bool = True,
@@ -192,6 +192,7 @@ class EmbeddingModelBase(ABC):
         """
         Evaluate the finetuned embedding model
         """
+
         payload: EmbeddingEvaluationRequest = {
             "datasource_name_or_id": datasource.id,
             "datasource_label_column": label_column,
@@ -219,17 +220,17 @@ class EmbeddingModelBase(ABC):
         else:
             raise ValueError("Invalid embedding model")
 
-        def get_result(task_id: str) -> ClassificationMetrics | RegressionMetrics:
+        def get_result(job_id: str) -> ClassificationMetrics | RegressionMetrics:
             client = OrcaClient._resolve_client()
             if isinstance(self, PretrainedEmbeddingModel):
                 res = client.GET(
-                    "/pretrained_embedding_model/{model_name}/evaluation/{task_id}",
-                    params={"model_name": self.name, "task_id": task_id},
+                    "/pretrained_embedding_model/{model_name}/evaluation/{job_id}",
+                    params={"model_name": self.name, "job_id": job_id},
                 )["result"]
             elif isinstance(self, FinetunedEmbeddingModel):
                 res = client.GET(
-                    "/finetuned_embedding_model/{name_or_id}/evaluation/{task_id}",
-                    params={"name_or_id": self.id, "task_id": task_id},
+                    "/finetuned_embedding_model/{name_or_id}/evaluation/{job_id}",
+                    params={"name_or_id": self.id, "job_id": job_id},
                 )["result"]
             else:
                 raise ValueError("Invalid embedding model")
@@ -263,7 +264,7 @@ class EmbeddingModelBase(ABC):
                 )
             )
 
-        job = Job(response["task_id"], lambda: get_result(response["task_id"]))
+        job = Job(response["job_id"], lambda: get_result(response["job_id"]))
         return job if background else job.result()
 
 
@@ -462,12 +463,13 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
     def finetune(
         self,
         name: str,
-        train_datasource: Datasource | LabeledMemoryset,
+        train_datasource: Datasource | LabeledMemoryset | ScoredMemoryset,
         *,
         eval_datasource: Datasource | None = None,
         label_column: str = "label",
+        score_column: str = "score",
         value_column: str = "value",
-        training_method: EmbeddingFinetuningMethod = "classification",
+        training_method: EmbeddingFinetuningMethod | None = None,
         training_args: dict | None = None,
         if_exists: CreateMode = "error",
         background: Literal[True],
@@ -478,12 +480,13 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
     def finetune(
         self,
         name: str,
-        train_datasource: Datasource | LabeledMemoryset,
+        train_datasource: Datasource | LabeledMemoryset | ScoredMemoryset,
         *,
         eval_datasource: Datasource | None = None,
         label_column: str = "label",
+        score_column: str = "score",
         value_column: str = "value",
-        training_method: EmbeddingFinetuningMethod = "classification",
+        training_method: EmbeddingFinetuningMethod | None = None,
         training_args: dict | None = None,
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
@@ -493,12 +496,13 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
     def finetune(
         self,
         name: str,
-        train_datasource: Datasource | LabeledMemoryset,
+        train_datasource: Datasource | LabeledMemoryset | ScoredMemoryset,
         *,
         eval_datasource: Datasource | None = None,
         label_column: str = "label",
+        score_column: str = "score",
         value_column: str = "value",
-        training_method: EmbeddingFinetuningMethod = "classification",
+        training_method: EmbeddingFinetuningMethod | None = None,
         training_args: dict | None = None,
         if_exists: CreateMode = "error",
         background: bool = False,
@@ -510,9 +514,10 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
             name: Name of the finetuned embedding model
             train_datasource: Data to train on
             eval_datasource: Optionally provide data to evaluate on
-            label_column: Column name of the label
+            label_column: Column name of the label.
+            score_column: Column name of the score (for regression when training on scored data).
             value_column: Column name of the value
-            training_method: Training method to use
+            training_method: Optional training method override. If omitted, Lighthouse defaults apply.
             training_args: Optional override for Hugging Face [`TrainingArguments`][transformers.TrainingArguments].
                 If not provided, reasonable training arguments will be used for the specified training method
             if_exists: What to do if a finetuned embedding model with the same name already exists, defaults to
@@ -543,19 +548,22 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
 
             return existing
 
-        from .memoryset import LabeledMemoryset
+        from .memoryset import LabeledMemoryset, ScoredMemoryset
 
         payload: FinetuneEmbeddingModelRequest = {
             "name": name,
             "base_model": self.name,
             "label_column": label_column,
+            "score_column": score_column,
             "value_column": value_column,
-            "training_method": training_method,
             "training_args": training_args or {},
         }
+        if training_method is not None:
+            payload["training_method"] = training_method
+
         if isinstance(train_datasource, Datasource):
             payload["train_datasource_name_or_id"] = train_datasource.id
-        elif isinstance(train_datasource, LabeledMemoryset):
+        elif isinstance(train_datasource, (LabeledMemoryset, ScoredMemoryset)):
             payload["train_memoryset_name_or_id"] = train_datasource.id
         if eval_datasource is not None:
             payload["eval_datasource_name_or_id"] = eval_datasource.id
@@ -566,7 +574,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
             json=payload,
         )
         job = Job(
-            res["finetuning_task_id"],
+            res["finetuning_job_id"],
             lambda: FinetunedEmbeddingModel.open(res["id"]),
         )
         return job if background else job.result()

orca_sdk/job.py

@@ -17,7 +17,7 @@ class JobConfig(TypedDict):
 
 
 class Status(Enum):
-    """Status of a cloud job in the task queue"""
+    """Status of a cloud job in the job queue"""
 
     # the INITIALIZED state should never be returned by the API
     INITIALIZED = "INITIALIZED"
@@ -141,8 +141,8 @@ class Job(Generic[TResult]):
             List of jobs matching the given filters
         """
         client = OrcaClient._resolve_client()
-        paginated_tasks = client.GET(
-            "/task",
+        paginated_jobs = client.GET(
+            "/job",
             params={
                 "status": (
                     [s.value for s in status]
@@ -175,7 +175,7 @@ class Job(Generic[TResult]):
                     obj,
                 )[-1]
             )(t)
-            for t in paginated_tasks["items"]
+            for t in paginated_jobs["items"]
         ]
 
     def __init__(self, id: str, get_value: Callable[[], TResult | None] | None = None):
@@ -188,29 +188,29 @@ class Job(Generic[TResult]):
         """
         self.id = id
         client = OrcaClient._resolve_client()
-        task = client.GET("/task/{task_id}", params={"task_id": id})
+        job = client.GET("/job/{job_id}", params={"job_id": id})
 
         def default_get_value():
             client = OrcaClient._resolve_client()
-            return cast(TResult | None, client.GET("/task/{task_id}", params={"task_id": id})["result"])
+            return cast(TResult | None, client.GET("/job/{job_id}", params={"job_id": id})["result"])
 
         self._get_value = get_value or default_get_value
-        self.type = task["type"]
-        self.status = Status(task["status"])
-        self.steps_total = task["steps_total"]
-        self.steps_completed = task["steps_completed"]
-        self.exception = task["exception"]
+        self.type = job["type"]
+        self.status = Status(job["status"])
+        self.steps_total = job["steps_total"]
+        self.steps_completed = job["steps_completed"]
+        self.exception = job["exception"]
         self.value = (
             None
-            if task["status"] != "COMPLETED"
+            if job["status"] != "COMPLETED"
             else (
                 get_value()
                 if get_value is not None
-                else cast(TResult, task["result"]) if task["result"] is not None else None
+                else cast(TResult, job["result"]) if job["result"] is not None else None
             )
         )
-        self.updated_at = datetime.fromisoformat(task["updated_at"])
-        self.created_at = datetime.fromisoformat(task["created_at"])
+        self.updated_at = datetime.fromisoformat(job["updated_at"])
+        self.created_at = datetime.fromisoformat(job["created_at"])
         self.refreshed_at = datetime.now()
 
     def refresh(self, throttle: float = 0):
@@ -227,7 +227,7 @@ class Job(Generic[TResult]):
         self.refreshed_at = current_time
 
         client = OrcaClient._resolve_client()
-        status_info = client.GET("/task/{task_id}/status", params={"task_id": self.id})
+        status_info = client.GET("/job/{job_id}/status", params={"job_id": self.id})
         self.status = Status(status_info["status"])
         if status_info["steps_total"] is not None:
             self.steps_total = status_info["steps_total"]
@@ -339,5 +339,5 @@ def abort(self, show_progress: bool = False, refresh_interval: int = 1, max_wait
         max_wait: Maximum time to wait for the job to abort in seconds
     """
     client = OrcaClient._resolve_client()
-    client.DELETE("/task/{task_id}/abort", params={"task_id": self.id})
+    client.DELETE("/job/{job_id}/abort", params={"job_id": self.id})
     self.wait(show_progress, refresh_interval, max_wait)