pixeltable 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

pixeltable/plan.py

@@ -76,7 +76,8 @@ class Analyzer:
                             f'order_by()'))
                     self.similarity_clause = similarity_clauses[0]
                     img_col = self.similarity_clause.img_col_ref.col
-                    if not img_col.is_indexed:
+                    indexed_col_ids = {info.col.id for info in tbl.tbl_version.idxs_by_name.values()}
+                    if img_col.id not in indexed_col_ids:
                         raise excs.Error(f'nearest() not available for unindexed column {img_col.name}')
 
         # all exprs that are evaluated in Python; not executable
@@ -220,18 +221,11 @@ class Planner:
     ) -> exec.ExecNode:
         """Creates a plan for TableVersion.insert()"""
         assert not tbl.is_view()
-        # things we need to materialize:
-        # 1. stored_cols: all cols we need to store, incl computed cols (and indices)
+        # stored_cols: all cols we need to store, incl computed cols (and indices)
         stored_cols = [c for c in tbl.cols if c.is_stored]
         assert len(stored_cols) > 0
-        # 2. values to insert into indices
-        indexed_cols = [c for c in tbl.cols if c.is_indexed]
-        index_info: List[Tuple[catalog.Column, func.Function]] = []
-        if len(indexed_cols) > 0:
-            from pixeltable.functions.nos.image_embedding import openai_clip
-            index_info = [(c, openai_clip) for c in tbl.cols if c.is_indexed]
 
-        row_builder = exprs.RowBuilder([], stored_cols, index_info, [])
+        row_builder = exprs.RowBuilder([], stored_cols, [])
 
         # create InMemoryDataNode for 'rows'
         stored_col_info = row_builder.output_slot_idxs()
@@ -260,7 +254,7 @@ class Planner:
     @classmethod
     def create_update_plan(
             cls, tbl: catalog.TableVersionPath,
-            update_targets: List[Tuple[catalog.Column, exprs.Expr]],
+            update_targets: dict[catalog.Column, exprs.Expr],
             recompute_targets: List[catalog.Column],
             where_clause: Optional[exprs.Predicate], cascade: bool
     ) -> Tuple[exec.ExecNode, List[str], List[catalog.Column]]:
@@ -279,7 +273,7 @@ class Planner:
         # retrieve all stored cols and all target exprs
         assert isinstance(tbl, catalog.TableVersionPath)
         target = tbl.tbl_version  # the one we need to update
-        updated_cols = [col for col, _ in update_targets]
+        updated_cols = list(update_targets.keys())
         if len(recompute_targets) > 0:
             recomputed_cols = recompute_targets.copy()
         else:
@@ -291,12 +285,12 @@ class Planner:
             col for col in target.cols if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
         ]
         select_list = [exprs.ColumnRef(col) for col in copied_cols]
-        select_list.extend([expr for _, expr in update_targets])
+        select_list.extend(update_targets.values())
 
         recomputed_exprs = \
             [c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols]
         # recomputed cols reference the new values of the updated cols
-        for col, e in update_targets:
+        for col, e in update_targets.items():
             exprs.Expr.list_substitute(recomputed_exprs, exprs.ColumnRef(col), e)
         select_list.extend(recomputed_exprs)
 
@@ -375,16 +369,10 @@ class Planner:
         #   the store
         target = view.tbl_version  # the one we need to populate
         stored_cols = [c for c in target.cols if c.is_stored and (c.is_computed or target.is_iterator_column(c))]
-        # 2. index values
-        indexed_cols = [c for c in target.cols if c.is_indexed]
-        index_info: List[Tuple[catalog.Column, func.Function]] = []
-        if len(indexed_cols) > 0:
-            from pixeltable.functions.nos.image_embedding import openai_clip
-            index_info = [(c, openai_clip) for c in target.cols if c.is_indexed]
-        # 3. for component views: iterator args
+        # 2. for component views: iterator args
         iterator_args = [target.iterator_args] if target.iterator_args is not None else []
 
-        row_builder = exprs.RowBuilder(iterator_args, stored_cols, index_info, [])
+        row_builder = exprs.RowBuilder(iterator_args, stored_cols, [])
 
         # execution plan:
         # 1. materialize exprs computed from the base that are needed for stored view columns
@@ -548,7 +536,7 @@ class Planner:
         analyzer = Analyzer(
             tbl, select_list, where_clause=where_clause, group_by_clause=group_by_clause,
             order_by_clause=order_by_clause)
-        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], [], analyzer.sql_exprs)
+        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], analyzer.sql_exprs)
 
         analyzer.finalize(row_builder)
         # select_list: we need to materialize everything that's been collected
@@ -627,21 +615,15 @@ class Planner:
     @classmethod
     def create_add_column_plan(
             cls, tbl: catalog.TableVersionPath, col: catalog.Column
-    ) -> Tuple[exec.ExecNode, Optional[int], Optional[int]]:
+    ) -> Tuple[exec.ExecNode, Optional[int]]:
         """Creates a plan for InsertableTable.add_column()
         Returns:
             plan: the plan to execute
-            ctx: the context to use for the plan
             value_expr slot idx for the plan output (for computed cols)
-            embedding slot idx for the plan output (for indexed image cols)
         """
         assert isinstance(tbl, catalog.TableVersionPath)
         index_info: List[Tuple[catalog.Column, func.Function]] = []
-        if col.is_indexed:
-            from pixeltable.functions.nos.image_embedding import openai_clip
-            index_info = [(col, openai_clip)]
-        row_builder = exprs.RowBuilder(
-            output_exprs=[], columns=[col], indices=index_info, input_exprs=[])
+        row_builder = exprs.RowBuilder(output_exprs=[], columns=[col], input_exprs=[])
         analyzer = Analyzer(tbl, row_builder.default_eval_ctx.target_exprs)
         plan = cls._create_query_plan(tbl, row_builder=row_builder, analyzer=analyzer, with_pk=True)
         plan.ctx.batch_size = 16
@@ -651,6 +633,5 @@ class Planner:
         # we want to flush images
         if col.is_computed and col.is_stored and col.col_type.is_image_type():
             plan.set_stored_img_cols(row_builder.output_slot_idxs())
-        value_expr_slot_idx: Optional[int] = row_builder.output_slot_idxs()[0].slot_idx if col.is_computed else None
-        embedding_slot_idx: Optional[int] = row_builder.index_slot_idxs()[0].slot_idx if col.is_indexed else None
-        return plan, value_expr_slot_idx, embedding_slot_idx
+        value_expr_slot_idx = row_builder.output_slot_idxs()[0].slot_idx if col.is_computed else None
+        return plan, value_expr_slot_idx

pixeltable/store.py

@@ -38,7 +38,7 @@ class StoreBase:
         self.tbl_version = tbl_version
         self.sa_md = sql.MetaData()
         self.sa_tbl: Optional[sql.Table] = None
-        self._create_sa_tbl()
+        self.create_sa_tbl()
 
     def pk_columns(self) -> List[sql.Column]:
         return self._pk_columns
@@ -62,7 +62,7 @@ class StoreBase:
         return [*rowid_cols, self.v_min_col, self.v_max_col]
 
 
-    def _create_sa_tbl(self) -> None:
+    def create_sa_tbl(self) -> None:
         """Create self.sa_tbl from self.tbl_version."""
         system_cols = self._create_system_columns()
         all_cols = system_cols.copy()
@@ -76,9 +76,6 @@ class StoreBase:
                 all_cols.append(col.sa_errormsg_col)
                 all_cols.append(col.sa_errortype_col)
 
-            if col.is_indexed:
-                all_cols.append(col.sa_idx_col)
-
             # we create an index for:
             # - scalar columns (except for strings, because long strings can't be used for B-tree indices)
             # - non-computed video and image columns (they will contain external paths/urls that users might want to
@@ -145,8 +142,8 @@ class StoreBase:
         """Move tmp media files that we generated to a permanent location"""
         for c in media_cols:
             for table_row in table_rows:
-                file_url = table_row[c.storage_name()]
-                table_row[c.storage_name()] = self._move_tmp_media_file(file_url, c, v_min)
+                file_url = table_row[c.store_name()]
+                table_row[c.store_name()] = self._move_tmp_media_file(file_url, c, v_min)
 
     def _create_table_row(
             self, input_row: exprs.DataRow, row_builder: exprs.RowBuilder, media_cols: List[catalog.Column],
@@ -168,16 +165,19 @@ class StoreBase:
 
         return table_row, num_excs
 
-    def count(self) -> None:
+    def count(self, conn: Optional[sql.engine.Connection] = None) -> int:
         """Return the number of rows visible in self.tbl_version"""
         stmt = sql.select(sql.func.count('*'))\
             .select_from(self.sa_tbl)\
             .where(self.v_min_col <= self.tbl_version.version)\
             .where(self.v_max_col > self.tbl_version.version)
-        with env.Env.get().engine.begin() as conn:
+        if conn is None:
+            with env.Env.get().engine.connect() as conn:
+                result = conn.execute(stmt).scalar_one()
+        else:
             result = conn.execute(stmt).scalar_one()
-            assert isinstance(result, int)
-            return result
+        assert isinstance(result, int)
+        return result
 
     def create(self, conn: sql.engine.Connection) -> None:
         self.sa_md.create_all(bind=conn)
@@ -193,38 +193,35 @@ class StoreBase:
         message).
         """
         assert col.is_stored
-        stmt = sql.text(f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.storage_name()} {col.col_type.to_sql()}')
+        col_type_str = col.get_sa_col_type().compile(dialect=conn.dialect)
+        stmt = sql.text(f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.store_name()} {col_type_str} NULL')
         log_stmt(_logger, stmt)
         conn.execute(stmt)
-        added_storage_cols = [col.storage_name()]
+        added_storage_cols = [col.store_name()]
         if col.records_errors:
             # we also need to create the errormsg and errortype storage cols
             stmt = (f'ALTER TABLE {self._storage_name()} '
-                    f'ADD COLUMN {col.errormsg_storage_name()} {StringType().to_sql()} DEFAULT NULL')
+                    f'ADD COLUMN {col.errormsg_store_name()} {StringType().to_sql()} DEFAULT NULL')
             conn.execute(sql.text(stmt))
             stmt = (f'ALTER TABLE {self._storage_name()} '
-                    f'ADD COLUMN {col.errortype_storage_name()} {StringType().to_sql()} DEFAULT NULL')
+                    f'ADD COLUMN {col.errortype_store_name()} {StringType().to_sql()} DEFAULT NULL')
             conn.execute(sql.text(stmt))
-        added_storage_cols.extend([col.errormsg_storage_name(), col.errortype_storage_name()])
-        self._create_sa_tbl()
+            added_storage_cols.extend([col.errormsg_store_name(), col.errortype_store_name()])
+        self.create_sa_tbl()
         _logger.info(f'Added columns {added_storage_cols} to storage table {self._storage_name()}')
 
-    def drop_column(self, col: Optional[catalog.Column] = None, conn: Optional[sql.engine.Connection] = None) -> None:
-        """Re-create self.sa_tbl and drop column, if one is given"""
-        if col is not None:
-            assert conn is not None
-            stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.storage_name()}'
+    def drop_column(self, col: catalog.Column, conn: sql.engine.Connection) -> None:
+        """Execute Alter Table Drop Column statement"""
+        stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.store_name()}'
+        conn.execute(sql.text(stmt))
+        if col.records_errors:
+            stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errormsg_store_name()}'
+            conn.execute(sql.text(stmt))
+            stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errortype_store_name()}'
             conn.execute(sql.text(stmt))
-            if col.records_errors:
-                stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errormsg_storage_name()}'
-                conn.execute(sql.text(stmt))
-                stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errortype_storage_name()}'
-                conn.execute(sql.text(stmt))
-        self._create_sa_tbl()
 
     def load_column(
-            self, col: catalog.Column, exec_plan: ExecNode, value_expr_slot_idx: int, embedding_slot_idx: int,
-            conn: sql.engine.Connection
+            self, col: catalog.Column, exec_plan: ExecNode, value_expr_slot_idx: int, conn: sql.engine.Connection
     ) -> int:
         """Update store column of a computed column with values produced by an execution plan
 
@@ -253,18 +250,11 @@ class StoreBase:
                             col.sa_errormsg_col: error_msg
                         }
                     else:
-                        val = result_row.get_stored_val(value_expr_slot_idx)
+                        val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
                         if col.col_type.is_media_type():
                             val = self._move_tmp_media_file(val, col, result_row.pk[-1])
                         values_dict = {col.sa_col: val}
 
-                if col.is_indexed:
-                    # TODO: deal with exceptions
-                    assert not result_row.has_exc(embedding_slot_idx)
-                    # don't use get_stored_val() here, we need to pass the ndarray
-                    embedding = result_row[embedding_slot_idx]
-                    values_dict[col.sa_index_col] = embedding
-
                 update_stmt = sql.update(self.sa_tbl).values(values_dict)
                 for pk_col, pk_val in zip(self.pk_columns(), result_row.pk):
                     update_stmt = update_stmt.where(pk_col == pk_val)
@@ -337,6 +327,7 @@ class StoreBase:
             self, current_version: int, base_versions: List[Optional[int]], match_on_vmin: bool,
             where_clause: Optional[sql.ClauseElement], conn: sql.engine.Connection) -> int:
         """Mark rows as deleted that are live and were created prior to current_version.
+        Also: populate the undo columns
         Args:
             base_versions: if non-None, join only to base rows that were created at that version,
                 otherwise join to rows that are live in the base's current version (which is distinct from the
@@ -354,8 +345,14 @@ class StoreBase:
         rowid_join_clause = self._rowid_join_predicate()
         base_versions_clause = sql.true() if len(base_versions) == 0 \
             else self.base._versions_clause(base_versions, match_on_vmin)
+        set_clause = {self.v_max_col: current_version}
+        for index_info in self.tbl_version.idxs_by_name.values():
+            # copy value column to undo column
+            set_clause[index_info.undo_col.sa_col] = index_info.val_col.sa_col
+            # set value column to NULL
+            set_clause[index_info.val_col.sa_col] = None
         stmt = sql.update(self.sa_tbl) \
-            .values({self.v_max_col: current_version}) \
+            .values(set_clause) \
             .where(where_clause) \
             .where(rowid_join_clause) \
             .where(base_versions_clause)
@@ -416,8 +413,8 @@ class StoreComponentView(StoreView):
         self.rowid_cols.append(self.pos_col)
         return self.rowid_cols
 
-    def _create_sa_tbl(self) -> None:
-        super()._create_sa_tbl()
+    def create_sa_tbl(self) -> None:
+        super().create_sa_tbl()
         # we need to fix up the 'pos' column in TableVersion
         self.tbl_version.cols_by_name['pos'].sa_col = self.pos_col
 

pixeltable/tests/conftest.py

@@ -6,11 +6,12 @@ from typing import List
 
 import numpy as np
 import pytest
+import PIL.Image
 
 import pixeltable as pxt
 import pixeltable.catalog as catalog
 from pixeltable import exprs
-from pixeltable import functions as ptf
+import pixeltable.functions as pxtf
 from pixeltable.exprs import RELATIVE_PATH_ROOT as R
 from pixeltable.metadata import SystemInfo, create_system_info
 from pixeltable.metadata.schema import TableSchemaVersion, TableVersion, Table, Function, Dir
@@ -80,8 +81,8 @@ def test_tbl(test_client: pxt.Client) -> catalog.Table:
     return create_test_tbl(test_client)
 
 # @pytest.fixture(scope='function')
-# def test_stored_fn(test_client: pt.Client) -> pt.Function:
-#     @pt.udf(return_type=pt.IntType(), param_types=[pt.IntType()])
+# def test_stored_fn(test_client: pxt.Client) -> pxt.Function:
+#     @pxt.udf(return_type=pxt.IntType(), param_types=[pxt.IntType()])
 #     def test_fn(x):
 #         return x + 1
 #     test_client.create_function('test_fn', test_fn)
@@ -89,7 +90,7 @@ def test_tbl(test_client: pxt.Client) -> catalog.Table:
 
 @pytest.fixture(scope='function')
 def test_tbl_exprs(test_tbl: catalog.Table) -> List[exprs.Expr]:
-#def test_tbl_exprs(test_tbl: catalog.Table, test_stored_fn: pt.Function) -> List[exprs.Expr]:
+#def test_tbl_exprs(test_tbl: catalog.Table, test_stored_fn: pxt.Function) -> List[exprs.Expr]:
 
     t = test_tbl
     return [
@@ -120,8 +121,7 @@ def test_tbl_exprs(test_tbl: catalog.Table) -> List[exprs.Expr]:
         t.c1.apply(json.loads),
         t.c8.errortype,
         t.c8.errormsg,
-        ptf.sum(t.c2, group_by=t.c4, order_by=t.c3),
-        #test_stored_fn(t.c2),
+        pxtf.sum(t.c2, group_by=t.c4, order_by=t.c3),
     ]
 
 @pytest.fixture(scope='function')
@@ -153,17 +153,11 @@ def img_tbl_exprs(img_tbl: catalog.Table) -> List[exprs.Expr]:
         img_t.img.localpath,
     ]
 
-# TODO: why does this not work with a session scope? (some user tables don't get created with create_all())
-#@pytest.fixture(scope='session')
-#def indexed_img_tbl(init_env: None) -> catalog.Table:
-#    cl = pt.Client()
-#    db = cl.create_db('test_indexed')
 @pytest.fixture(scope='function')
-def indexed_img_tbl(test_client: pxt.Client) -> catalog.Table:
-    skip_test_if_not_installed('nos')
+def small_img_tbl(test_client: pxt.Client) -> catalog.Table:
     cl = test_client
     schema = {
-        'img': { 'type': ImageType(nullable=False), 'indexed': True },
+        'img': ImageType(nullable=False),
         'category': StringType(nullable=False),
         'split': StringType(nullable=False),
     }

pixeltable/tests/ext/test_yolox.py

@@ -0,0 +1,21 @@
+import pixeltable as pxt
+from pixeltable.tests.utils import skip_test_if_not_installed, get_image_files, validate_update_status
+
+
+class TestYolox:
+
+    def test_yolox(self, test_client: pxt.Client):
+        skip_test_if_not_installed('yolox')
+        from pixeltable.ext.functions.yolox import yolox
+        cl = test_client
+        t = cl.create_table('yolox_test', {'image': pxt.ImageType()})
+        t['detect_yolox_tiny'] = yolox(t.image, model_id='yolox_tiny')
+        t['detect_yolox_nano'] = yolox(t.image, model_id='yolox_nano', threshold=0.2)
+        t['yolox_nano_bboxes'] = t.detect_yolox_nano.bboxes
+        images = get_image_files()[:10]
+        validate_update_status(t.insert({'image': image} for image in images), expected_rows=10)
+        rows = t.collect()
+        # Verify correctly formed JSON
+        assert all(list(result.keys()) == ['bboxes', 'labels', 'scores'] for result in rows['detect_yolox_tiny'])
+        # Verify that bboxes are actually present in at least some of the rows.
+        assert any(len(bboxes) > 0 for bboxes in rows['yolox_nano_bboxes'])

pixeltable/tests/functions/test_fireworks.py

@@ -0,0 +1,43 @@
+import pytest
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.tests.utils import skip_test_if_not_installed, validate_update_status
+
+
+@pytest.mark.remote_api
+class TestFireworks:
+
+    def test_fireworks(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('fireworks')
+        TestFireworks.skip_test_if_no_fireworks_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': pxt.StringType()})
+        from pixeltable.functions.fireworks import chat_completions
+        messages = [{'role': 'user', 'content': t.input}]
+        t['output'] = chat_completions(
+            messages=messages,
+            model='accounts/fireworks/models/llama-v2-7b-chat'
+        )
+        t['output_2'] = chat_completions(
+            messages=messages,
+            model='accounts/fireworks/models/llama-v2-7b-chat',
+            max_tokens=300,
+            top_k=40,
+            top_p=0.9,
+            temperature=0.7
+        )
+        validate_update_status(t.insert(input="How's everything going today?"), 1)
+        results = t.collect()
+        assert len(results['output'][0]['choices'][0]['message']['content']) > 0
+        assert len(results['output_2'][0]['choices'][0]['message']['content']) > 0
+
+    # This ensures that the test will be skipped, rather than returning an error, when no API key is
+    # available (for example, when a PR runs in CI).
+    @staticmethod
+    def skip_test_if_no_fireworks_client() -> None:
+        try:
+            import pixeltable.functions.fireworks
+            _ = pixeltable.functions.fireworks.fireworks_client()
+        except excs.Error as exc:
+            pytest.skip(str(exc))

pixeltable/tests/functions/test_functions.py

@@ -0,0 +1,60 @@
+import pixeltable as pxt
+from pixeltable import catalog
+from pixeltable.functions.pil.image import blend
+from pixeltable.iterators import FrameIterator
+from pixeltable.tests.utils import get_video_files, skip_test_if_not_installed
+from pixeltable.type_system import VideoType, StringType
+
+
+class TestFunctions:
+    def test_pil(self, img_tbl: catalog.Table) -> None:
+        t = img_tbl
+        _ = t[t.img, t.img.rotate(90), blend(t.img, t.img.rotate(90), 0.5)].show()
+
+    def test_eval_detections(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('nos')
+        cl = test_client
+        video_t = cl.create_table('video_tbl', {'video': VideoType()})
+        # create frame view
+        args = {'video': video_t.video, 'fps': 1}
+        v = cl.create_view('test_view', video_t, iterator_class=FrameIterator, iterator_args=args)
+
+        files = get_video_files()
+        video_t.insert(video=files[-1])
+        v.add_column(frame_s=v.frame.resize([640, 480]))
+        from pixeltable.functions.nos.object_detection_2d import yolox_nano, yolox_small, yolox_large
+        v.add_column(detections_a=yolox_nano(v.frame_s))
+        v.add_column(detections_b=yolox_small(v.frame_s))
+        v.add_column(gt=yolox_large(v.frame_s))
+        from pixeltable.functions.eval import eval_detections, mean_ap
+        res = v.select(
+            eval_detections(
+                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels
+            )).show()
+        v.add_column(
+            eval_a=eval_detections(
+                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels))
+        v.add_column(
+            eval_b=eval_detections(
+                v.detections_b.bboxes, v.detections_b.labels, v.detections_b.scores, v.gt.bboxes, v.gt.labels))
+        ap_a = v.select(mean_ap(v.eval_a)).show()[0, 0]
+        ap_b = v.select(mean_ap(v.eval_b)).show()[0, 0]
+        common_classes = set(ap_a.keys()) & set(ap_b.keys())
+
+        ## TODO: following assertion is failing on CI, 
+        # It is not necessarily a bug, as assert codition is not expected to be always true
+        # for k in common_classes:
+        # assert ap_a[k] <= ap_b[k]
+
+    def test_str(self, test_client: pxt.Client) -> None:
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.string import str_format
+        t.add_column(s1=str_format('ABC {0}', t.input))
+        t.add_column(s2=str_format('DEF {this}', this=t.input))
+        t.add_column(s3=str_format('GHI {0} JKL {this}', t.input, this=t.input))
+        status = t.insert(input='MNO')
+        assert status.num_rows == 1
+        assert status.num_excs == 0
+        row = t.head()[0]
+        assert row == {'input': 'MNO', 's1': 'ABC MNO', 's2': 'DEF MNO', 's3': 'GHI MNO JKL MNO'}

pixeltable/tests/{test_functions.py → functions/test_huggingface.py}

@@ -3,144 +3,12 @@ from typing import Dict, Any
 import pytest
 
 import pixeltable as pxt
-from pixeltable import catalog
-from pixeltable.env import Env
-import pixeltable.exceptions as excs
-from pixeltable.functions.pil.image import blend
-from pixeltable.iterators import FrameIterator
-from pixeltable.tests.utils import get_video_files, skip_test_if_not_installed, get_sentences, get_image_files
-from pixeltable.type_system import VideoType, StringType, JsonType, ImageType, BoolType, FloatType, ArrayType
+from pixeltable.tests.utils import skip_test_if_not_installed, get_sentences, get_image_files, \
+    SAMPLE_IMAGE_URL
+from pixeltable.type_system import StringType, JsonType, ImageType, BoolType, FloatType, ArrayType
 
 
-class TestFunctions:
-    def test_pil(self, img_tbl: catalog.Table) -> None:
-        t = img_tbl
-        _ = t[t.img, t.img.rotate(90), blend(t.img, t.img.rotate(90), 0.5)].show()
-
-    def test_eval_detections(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('nos')
-        cl = test_client
-        video_t = cl.create_table('video_tbl', {'video': VideoType()})
-        # create frame view
-        args = {'video': video_t.video, 'fps': 1}
-        v = cl.create_view('test_view', video_t, iterator_class=FrameIterator, iterator_args=args)
-
-        files = get_video_files()
-        video_t.insert(video=files[-1])
-        v.add_column(frame_s=v.frame.resize([640, 480]))
-        from pixeltable.functions.nos.object_detection_2d import yolox_nano, yolox_small, yolox_large
-        v.add_column(detections_a=yolox_nano(v.frame_s))
-        v.add_column(detections_b=yolox_small(v.frame_s))
-        v.add_column(gt=yolox_large(v.frame_s))
-        from pixeltable.functions.eval import eval_detections, mean_ap
-        res = v.select(
-            eval_detections(
-                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels
-            )).show()
-        v.add_column(
-            eval_a=eval_detections(
-                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels))
-        v.add_column(
-            eval_b=eval_detections(
-                v.detections_b.bboxes, v.detections_b.labels, v.detections_b.scores, v.gt.bboxes, v.gt.labels))
-        ap_a = v.select(mean_ap(v.eval_a)).show()[0, 0]
-        ap_b = v.select(mean_ap(v.eval_b)).show()[0, 0]
-        common_classes = set(ap_a.keys()) & set(ap_b.keys())
-
-        ## TODO: following assertion is failing on CI, 
-        # It is not necessarily a bug, as assert codition is not expected to be always true
-        # for k in common_classes:
-        # assert ap_a[k] <= ap_b[k]
-
-    def test_str(self, test_client: pxt.Client) -> None:
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.string import str_format
-        t.add_column(s1=str_format('ABC {0}', t.input))
-        t.add_column(s2=str_format('DEF {this}', this=t.input))
-        t.add_column(s3=str_format('GHI {0} JKL {this}', t.input, this=t.input))
-        status = t.insert(input='MNO')
-        assert status.num_rows == 1
-        assert status.num_excs == 0
-        row = t.head()[0]
-        assert row == {'input': 'MNO', 's1': 'ABC MNO', 's2': 'DEF MNO', 's3': 'GHI MNO JKL MNO'}
-
-    def test_openai(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('openai')
-        TestFunctions.skip_test_if_no_openai_client()
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.openai import chat_completions, embeddings, moderations
-        msgs = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": t.input}
-        ]
-        t.add_column(input_msgs=msgs)
-        t.add_column(chat_output=chat_completions(model='gpt-3.5-turbo', messages=t.input_msgs))
-        # with inlined messages
-        t.add_column(chat_output2=chat_completions(model='gpt-3.5-turbo', messages=msgs))
-        t.add_column(ada_embed=embeddings(model='text-embedding-ada-002', input=t.input))
-        t.add_column(text_3=embeddings(model='text-embedding-3-small', input=t.input))
-        t.add_column(moderation=moderations(input=t.input))
-        t.insert(input='I find you really annoying')
-        _ = t.head()
-
-    def test_gpt_4_vision(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('openai')
-        TestFunctions.skip_test_if_no_openai_client()
-        cl = test_client
-        t = cl.create_table('test_tbl', {'prompt': StringType(), 'img': ImageType()})
-        from pixeltable.functions.openai import chat_completions
-        from pixeltable.functions.string import str_format
-        msgs = [
-            {'role': 'user',
-             'content': [
-                 {'type': 'text', 'text': t.prompt},
-                 {'type': 'image_url', 'image_url': {
-                     'url': str_format('data:image/png;base64,{0}', t.img.b64_encode())
-                 }}
-             ]}
-        ]
-        t.add_column(response=chat_completions(model='gpt-4-vision-preview', messages=msgs, max_tokens=300))
-        t.add_column(response_content=t.response.choices[0].message.content)
-        t.insert(prompt="What's in this image?", img=_sample_image_url)
-        result = t.collect()['response_content'][0]
-        assert len(result) > 0
-
-    @staticmethod
-    def skip_test_if_no_openai_client() -> None:
-        try:
-            _ = Env.get().openai_client
-        except excs.Error as exc:
-            pytest.skip(str(exc))
-
-    def test_together(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('together')
-        if not Env.get().has_together_client:
-            pytest.skip(f'Together client does not exist (missing API key?)')
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.together import completions
-        t.add_column(output=completions(prompt=t.input, model='mistralai/Mixtral-8x7B-v0.1', stop=['\n']))
-        t.add_column(output_text=t.output.output.choices[0].text)
-        t.insert(input='I am going to the ')
-        result = t.select(t.output_text).collect()['output_text'][0]
-        assert len(result) > 0
-
-    def test_fireworks(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('fireworks')
-        try:
-            from pixeltable.functions.fireworks import initialize
-            initialize()
-        except:
-            pytest.skip(f'Fireworks client does not exist (missing API key?)')
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.fireworks import chat_completions
-        t['output'] = chat_completions(prompt=t.input, model='accounts/fireworks/models/llama-v2-7b-chat', max_tokens=256).choices[0].text
-        t.insert(input='I am going to the ')
-        result = t.select(t.output).collect()['output'][0]
-        assert len(result) > 0
+class TestHuggingface:
 
     def test_hf_function(self, test_client: pxt.Client) -> None:
         skip_test_if_not_installed('sentence_transformers')
@@ -255,10 +123,10 @@ class TestFunctions:
         for idx, model_id in enumerate(model_ids):
             col_name = f'embed_text{idx}'
             t[col_name] = clip_text(t.text, model_id=model_id)
-            assert t.column_types()[col_name] == ArrayType((None,), dtype=FloatType(), nullable=False)
+            assert t.column_types()[col_name].is_array_type()
             col_name = f'embed_img{idx}'
             t[col_name] = clip_image(t.img, model_id=model_id)
-            assert t.column_types()[col_name] == ArrayType((None,), dtype=FloatType(), nullable=False)
+            assert t.column_types()[col_name].is_array_type()
 
         def verify_row(row: Dict[str, Any]) -> None:
             for idx, _ in enumerate(model_ids):
@@ -281,14 +149,10 @@ class TestFunctions:
         t = cl.create_table('test_tbl', {'img': ImageType()})
         from pixeltable.functions.huggingface import detr_for_object_detection
         t['detect'] = detr_for_object_detection(t.img, model_id='facebook/detr-resnet-50', threshold=0.8)
-        status = t.insert(img=_sample_image_url)
+        status = t.insert(img=SAMPLE_IMAGE_URL)
         assert status.num_rows == 1
         assert status.num_excs == 0
         result = t.select(t.detect).collect()[0]['detect']
         assert 'orange' in result['label_text']
         assert 'bowl' in result['label_text']
         assert 'broccoli' in result['label_text']
-
-
-_sample_image_url = \
-    'https://raw.githubusercontent.com/pixeltable/pixeltable/master/docs/source/data/images/000000000009.jpg'