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

pixeltable/functions/huggingface.py

@@ -1,4 +1,4 @@
-from typing import Callable, TypeVar, Optional
+from typing import Callable, TypeVar, Optional, Any
 
 import PIL.Image
 import numpy as np
@@ -14,6 +14,7 @@ from pixeltable.functions.util import resolve_torch_device
 def sentence_transformer(
         sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False
 ) -> Batch[np.ndarray]:
+    """Runs the specified sentence transformer model."""
     env.Env.get().require_package('sentence_transformers')
     from sentence_transformers import SentenceTransformer
 
@@ -23,6 +24,16 @@ def sentence_transformer(
     return [array[i] for i in range(array.shape[0])]
 
 
+@sentence_transformer.conditional_return_type
+def _(model_id: str) -> ts.ArrayType:
+    try:
+        from sentence_transformers import SentenceTransformer
+        model = _lookup_model(model_id, SentenceTransformer)
+        return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
+    except ImportError:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+
+
 @pxt.udf
 def sentence_transformer_list(sentences: list, *, model_id: str, normalize_embeddings: bool = False) -> list:
     env.Env.get().require_package('sentence_transformers')
@@ -36,6 +47,7 @@ def sentence_transformer_list(sentences: list, *, model_id: str, normalize_embed
 
 @pxt.udf(batch_size=32)
 def cross_encoder(sentences1: Batch[str], sentences2: Batch[str], *, model_id: str) -> Batch[float]:
+    """Runs the specified cross-encoder model."""
     env.Env.get().require_package('sentence_transformers')
     from sentence_transformers import CrossEncoder
 
@@ -56,15 +68,15 @@ def cross_encoder_list(sentence1: str, sentences2: list, *, model_id: str) -> li
     return array.tolist()
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((512,), dtype=ts.FloatType(), nullable=False))
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
 def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
+    """Runs the specified CLIP model on text."""
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
     import torch
     from transformers import CLIPModel, CLIPProcessor
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
-    assert model.config.projection_dim == 512
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
     with torch.no_grad():
@@ -74,15 +86,15 @@ def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((512,), dtype=ts.FloatType(), nullable=False))
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
 def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndarray]:
+    """Runs the specified CLIP model on images."""
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
     import torch
     from transformers import CLIPModel, CLIPProcessor
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
-    assert model.config.projection_dim == 512
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
     with torch.no_grad():
@@ -92,8 +104,20 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
+@clip_text.conditional_return_type
+@clip_image.conditional_return_type
+def _(model_id: str) -> ts.ArrayType:
+    try:
+        from transformers import CLIPModel
+        model = _lookup_model(model_id, CLIPModel.from_pretrained)
+        return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
+    except ImportError:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+
+
 @pxt.udf(batch_size=4)
 def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
+    """Runs the specified DETR model."""
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
     import torch
@@ -121,6 +145,25 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     ]
 
 
+@pxt.udf
+def detr_to_coco(image: PIL.Image.Image, detr_info: dict[str, Any]) -> dict[str, Any]:
+    bboxes, labels = detr_info['boxes'], detr_info['labels']
+    annotations = [
+        {
+            'bbox': [bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]],
+            'category': label
+        }
+        for bbox, label in zip(bboxes, labels)
+    ]
+    return {
+        'image': {
+            'width': image.width,
+            'height': image.height
+        },
+        'annotations': annotations
+    }
+
+
 T = TypeVar('T')
 
 

pixeltable/functions/openai.py

@@ -16,8 +16,13 @@ from pixeltable import env
 from pixeltable.func import Batch
 
 
-def openai_client() -> openai.OpenAI:
-    return env.Env.get().get_client('openai', lambda api_key: openai.OpenAI(api_key=api_key))
+@env.register_client('openai')
+def _(api_key: str) -> openai.OpenAI:
+    return openai.OpenAI(api_key=api_key)
+
+
+def _openai_client() -> openai.OpenAI:
+    return env.Env.get().get_client('openai')
 
 
 # Exponential backoff decorator using tenacity.
@@ -44,7 +49,7 @@ def speech(
         response_format: Optional[str] = None,
         speed: Optional[float] = None
 ) -> str:
-    content = openai_client().audio.speech.create(
+    content = _openai_client().audio.speech.create(
         input=input,
         model=model,
         voice=voice,
@@ -53,7 +58,7 @@ def speech(
     )
     ext = response_format or 'mp3'
     output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
-    content.stream_to_file(output_filename, chunk_size=1 << 20)
+    content.write_to_file(output_filename)
     return output_filename
 
 
@@ -71,7 +76,7 @@ def transcriptions(
         temperature: Optional[float] = None
 ) -> dict:
     file = pathlib.Path(audio)
-    transcription = openai_client().audio.transcriptions.create(
+    transcription = _openai_client().audio.transcriptions.create(
         file=file,
         model=model,
         language=_opt(language),
@@ -93,7 +98,7 @@ def translations(
         temperature: Optional[float] = None
 ) -> dict:
     file = pathlib.Path(audio)
-    translation = openai_client().audio.translations.create(
+    translation = _openai_client().audio.translations.create(
         file=file,
         model=model,
         prompt=_opt(prompt),
@@ -127,7 +132,7 @@ def chat_completions(
         tool_choice: Optional[dict] = None,
         user: Optional[str] = None
 ) -> dict:
-    result = openai_client().chat.completions.create(
+    result = _openai_client().chat.completions.create(
         messages=messages,
         model=model,
         frequency_penalty=_opt(frequency_penalty),
@@ -171,7 +176,7 @@ def vision(
              }}
          ]}
     ]
-    result = openai_client().chat.completions.create(
+    result = _openai_client().chat.completions.create(
         messages=messages,
         model=model
     )
@@ -181,17 +186,26 @@ def vision(
 #####################################
 # Embeddings Endpoints
 
+_embedding_dimensions_cache: dict[str, int] = {
+    'text-embedding-ada-002': 1536,
+    'text-embedding-3-small': 1536,
+    'text-embedding-3-large': 3072,
+}
+
+
 @pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
 @_retry
 def embeddings(
         input: Batch[str],
         *,
         model: str,
+        dimensions: Optional[int] = None,
         user: Optional[str] = None
 ) -> Batch[np.ndarray]:
-    result = openai_client().embeddings.create(
+    result = _openai_client().embeddings.create(
         input=input,
         model=model,
+        dimensions=_opt(dimensions),
         user=_opt(user),
         encoding_format='float'
     )
@@ -201,6 +215,16 @@ def embeddings(
     ]
 
 
+@embeddings.conditional_return_type
+def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
+    if dimensions is None:
+        if model not in _embedding_dimensions_cache:
+            # TODO: find some other way to retrieve a sample
+            return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+        dimensions = _embedding_dimensions_cache.get(model, None)
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
+
+
 #####################################
 # Images Endpoints
 
@@ -216,7 +240,7 @@ def image_generations(
         user: Optional[str] = None
 ) -> PIL.Image.Image:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
-    result = openai_client().images.generate(
+    result = _openai_client().images.generate(
         prompt=prompt,
         model=_opt(model),
         quality=_opt(quality),
@@ -232,6 +256,20 @@ def image_generations(
     return img
 
 
+@image_generations.conditional_return_type
+def _(size: Optional[str] = None) -> ts.ImageType:
+    if size is None:
+        return ts.ImageType(size=(1024, 1024))
+    x_pos = size.find('x')
+    if x_pos == -1:
+        return ts.ImageType()
+    try:
+        width, height = int(size[:x_pos]), int(size[x_pos + 1:])
+    except ValueError:
+        return ts.ImageType()
+    return ts.ImageType(size=(width, height))
+
+
 #####################################
 # Moderations Endpoints
 
@@ -242,7 +280,7 @@ def moderations(
         *,
         model: Optional[str] = None
 ) -> dict:
-    result = openai_client().moderations.create(
+    result = _openai_client().moderations.create(
         input=input,
         model=_opt(model)
     )

pixeltable/functions/pil/image.py

@@ -1,16 +1,12 @@
-from typing import Dict, Any, Tuple, Optional
+from typing import Tuple, Optional
 
 import PIL.Image
+from PIL.Image import Dither
 
-from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType, BoolType
 import pixeltable.func as func
+from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType
 
 
-def _caller_return_type(bound_args: Optional[Dict[str, Any]]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    return bound_args['self'].col_type
-
 @func.udf(
     py_fn=PIL.Image.alpha_composite, return_type=ImageType(), param_types=[ImageType(), ImageType()])
 def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
@@ -28,71 +24,78 @@ def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.
 # PIL.Image.Image methods
 
 # Image.convert()
-def _convert_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    assert 'self' in bound_args
-    assert 'mode' in bound_args
-    img_type = bound_args['self'].col_type
-    return ImageType(size=img_type.size, mode=bound_args['mode'])
-@func.udf(return_type=_convert_return_type, param_types=[ImageType(), StringType()])
+@func.udf(param_types=[ImageType(), StringType()])
 def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
     return self.convert(mode)
 
+
+@convert.conditional_return_type
+def _(self: PIL.Image.Image, mode: str) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+
+
 # Image.crop()
-def _crop_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    img_type = bound_args['self'].col_type
-    box = bound_args['box']
-    if isinstance(box, list) and all(isinstance(x, int) for x in box):
-        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=img_type.mode)
-    return ImageType()  # we can't compute the size statically
 @func.udf(
-    py_fn=PIL.Image.Image.crop, return_type=_crop_return_type,
+    py_fn=PIL.Image.Image.crop,
     param_types=[ImageType(), ArrayType((4,), dtype=IntType())])
 def crop(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> PIL.Image.Image:
     pass
 
+@crop.conditional_return_type
+def _(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    if isinstance(box, list) and all(isinstance(x, int) for x in box):
+        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+
 # Image.getchannel()
-def _getchannel_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    img_type = bound_args['self'].col_type
-    return ImageType(size=img_type.size, mode='L')
-@func.udf(
-    py_fn=PIL.Image.Image.getchannel, return_type=_getchannel_return_type, param_types=[ImageType(), IntType()])
+@func.udf(py_fn=PIL.Image.Image.getchannel, param_types=[ImageType(), IntType()])
 def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
     pass
 
+@getchannel.conditional_return_type
+def _(self: PIL.Image.Image) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
+
+
 # Image.resize()
-def resize_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    assert 'size' in bound_args
-    return ImageType(size=bound_args['size'])
-@func.udf(return_type=resize_return_type, param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
+@func.udf(param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
 def resize(self: PIL.Image.Image, size: Tuple[int, int]) -> PIL.Image.Image:
     return self.resize(size)
 
+@resize.conditional_return_type
+def _(self: PIL.Image.Image, size: Tuple[int, int]) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+
 # Image.rotate()
-@func.udf(return_type=ImageType(), param_types=[ImageType(), IntType()])
+@func.udf(param_types=[ImageType(), IntType()])
 def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
     return self.rotate(angle)
 
-# Image.transform()
-@func.udf(return_type= _caller_return_type, param_types=[ImageType(), ArrayType((2,), dtype=IntType()), IntType()])
-def transform(self: PIL.Image.Image, size: Tuple[int, int], method: int) -> PIL.Image.Image:
-    return self.transform(size, method)
+@func.udf(py_fn=PIL.Image.Image.effect_spread, param_types=[ImageType(), IntType()])
+def effect_spread(self: PIL.Image.Image, distance: int) -> PIL.Image.Image:
+    pass
 
-@func.udf(
-    py_fn=PIL.Image.Image.effect_spread, return_type=_caller_return_type, param_types=[ImageType(), FloatType()])
-def effect_spread(self: PIL.Image.Image, distance: float) -> PIL.Image.Image:
+@func.udf(py_fn=PIL.Image.Image.transpose, param_types=[ImageType(), IntType()])
+def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
     pass
 
+@rotate.conditional_return_type
+@effect_spread.conditional_return_type
+@transpose.conditional_return_type
+def _(self: PIL.Image.Image) -> ColumnType:
+    return self.col_type
+
 @func.udf(
     py_fn=PIL.Image.Image.entropy, return_type=FloatType(), param_types=[ImageType(), ImageType(), JsonType()])
-def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> float:
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> float:
     pass
 
 @func.udf(py_fn=PIL.Image.Image.getbands, return_type=JsonType(), param_types=[ImageType()])
@@ -103,8 +106,7 @@ def getbands(self: PIL.Image.Image) -> Tuple[str]:
 def getbbox(self: PIL.Image.Image) -> Tuple[int, int, int, int]:
     pass
 
-@func.udf(
-    py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
+@func.udf(py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
 def getcolors(self: PIL.Image.Image, maxcolors: int) -> Tuple[Tuple[int, int, int], int]:
     pass
 
@@ -114,37 +116,32 @@ def getextrema(self: PIL.Image.Image) -> Tuple[int, int]:
 
 @func.udf(
     py_fn=PIL.Image.Image.getpalette, return_type=JsonType(), param_types=[ImageType(), StringType()])
-def getpalette(self: PIL.Image.Image, mode: str) -> Tuple[int]:
+def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> Tuple[int]:
     pass
 
 @func.udf(
-    py_fn=PIL.Image.Image.getpixel, return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
-def getpixel(self: PIL.Image.Image, xy: Tuple[int, int]) -> Tuple[int]:
-    pass
+    return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
+def getpixel(self: PIL.Image.Image, xy: tuple[int, int]) -> Tuple[int]:
+    # `xy` will be a list; `tuple(xy)` is necessary for pillow 9 compatibility
+    return self.getpixel(tuple(xy))
 
-@func.udf(
-    py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
+@func.udf(py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
 def getprojection(self: PIL.Image.Image) -> Tuple[int]:
     pass
 
-@func.udf(
-    py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
-def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> Tuple[int]:
+@func.udf(py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> Tuple[int]:
     pass
 
 @func.udf(
     py_fn=PIL.Image.Image.quantize, return_type=ImageType(),
     param_types=[ImageType(), IntType(), IntType(nullable=True), IntType(), IntType(nullable=True), IntType()])
 def quantize(
-        self: PIL.Image.Image, colors: int, method: int, kmeans: int, palette: int, dither: int) -> PIL.Image.Image:
+        self: PIL.Image.Image, colors: int = 256, method: Optional[int] = None, kmeans: int = 0,
+        palette: Optional[int] = None, dither: int = Dither.FLOYDSTEINBERG) -> PIL.Image.Image:
     pass
 
 @func.udf(
     py_fn=PIL.Image.Image.reduce, return_type=ImageType(), param_types=[ImageType(), IntType(), JsonType()])
-def reduce(self: PIL.Image.Image, factor: int, filter: Tuple[int]) -> PIL.Image.Image:
-    pass
-
-@func.udf(
-    py_fn=PIL.Image.Image.transpose, return_type=_caller_return_type, param_types=[ImageType(), IntType()])
-def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[Tuple[int]]) -> PIL.Image.Image:
     pass

pixeltable/functions/together.py

@@ -11,8 +11,13 @@ from pixeltable import env
 from pixeltable.func import Batch
 
 
-def together_client() -> together.Together:
-    return env.Env.get().get_client('together', lambda api_key: together.Together(api_key=api_key))
+@env.register_client('together')
+def _(api_key: str) -> together.Together:
+    return together.Together(api_key=api_key)
+
+
+def _together_client() -> together.Together:
+    return env.Env.get().get_client('together')
 
 
 @pxt.udf
@@ -31,7 +36,7 @@ def completions(
         n: Optional[int] = None,
         safety_model: Optional[str] = None
 ) -> dict:
-    return together_client().completions.create(
+    return _together_client().completions.create(
         prompt=prompt,
         model=model,
         max_tokens=max_tokens,
@@ -66,7 +71,7 @@ def chat_completions(
         tools: Optional[dict] = None,
         tool_choice: Optional[dict] = None
 ) -> dict:
-    return together_client().chat.completions.create(
+    return _together_client().chat.completions.create(
         messages=messages,
         model=model,
         max_tokens=max_tokens,
@@ -85,15 +90,36 @@ def chat_completions(
     ).dict()
 
 
+_embedding_dimensions_cache = {
+    'togethercomputer/m2-bert-80M-2k-retrieval': 768,
+    'togethercomputer/m2-bert-80M-8k-retrieval': 768,
+    'togethercomputer/m2-bert-80M-32k-retrieval': 768,
+    'WhereIsAI/UAE-Large-V1': 1024,
+    'BAAI/bge-large-en-v1.5': 1024,
+    'BAAI/bge-base-en-v1.5': 768,
+    'sentence-transformers/msmarco-bert-base-dot-v5': 768,
+    'bert-base-uncased': 768,
+}
+
+
 @pxt.udf(batch_size=32, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
 def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
-    result = together_client().embeddings.create(input=input, model=model)
+    result = _together_client().embeddings.create(input=input, model=model)
     return [
         np.array(data.embedding, dtype=np.float64)
         for data in result.data
     ]
 
 
+@embeddings.conditional_return_type
+def _(model: str) -> pxt.ArrayType:
+    if model not in _embedding_dimensions_cache:
+        # TODO: find some other way to retrieve a sample
+        return pxt.ArrayType((None,), dtype=pxt.FloatType())
+    dimensions = _embedding_dimensions_cache[model]
+    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+
+
 @pxt.udf
 def image_generations(
         prompt: str,
@@ -106,7 +132,7 @@ def image_generations(
         negative_prompt: Optional[str] = None,
 ) -> PIL.Image.Image:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
-    result = together_client().images.generate(
+    result = _together_client().images.generate(
         prompt=prompt,
         model=model,
         steps=steps,

pixeltable/functions/util.py

@@ -1,46 +1,3 @@
-from typing import Tuple, List, Optional
-import types
-import sys
-
-import pixeltable.func as func
-import pixeltable.type_system as ts
-import pixeltable.env as env
-
-
-def create_nos_modules() -> List[types.ModuleType]:
-    """Create module pixeltable.functions.nos with one submodule per task and return the submodules"""
-    models = env.Env.get().nos_client.ListModels()
-    model_info = [env.Env.get().nos_client.GetModelInfo(model) for model in models]
-    model_info.sort(key=lambda info: info.task.value)
-
-    module_name = 'pixeltable.functions.nos'
-    nos_module = types.ModuleType(module_name)
-    nos_module.__package__ = 'pixeltable.functions'
-    sys.modules[module_name] = nos_module
-
-    prev_task = ''
-    new_modules: List[types.ModuleType] = []
-    sub_module: Optional[types.ModuleType] = None
-    for info in model_info:
-        if info.task.value != prev_task:
-            # we construct one submodule per task
-            namespace = info.task.name.lower()
-            submodule_name = f'{module_name}.{namespace}'
-            sub_module = types.ModuleType(submodule_name)
-            sub_module.__package__ = module_name
-            setattr(nos_module, namespace, sub_module)
-            new_modules.append(sub_module)
-            sys.modules[submodule_name] = sub_module
-            prev_task = info.task.value
-
-        # add a Function for this model to the module
-        model_id = info.name.replace("/", "_").replace("-", "_")
-        pt_func = func.NOSFunction(info, f'{submodule_name}.{model_id}')
-        setattr(sub_module, model_id, pt_func)
-
-    return new_modules
-
-
 def resolve_torch_device(device: str) -> str:
     import torch
     if device == 'auto':

pixeltable/functions/video.py

@@ -1,14 +1,13 @@
-from typing import Optional
 import uuid
+from typing import Optional
+
 import av
-import sys
 
 import pixeltable.env as env
 import pixeltable.func as func
 import pixeltable.type_system as ts
 
-
-_format_defaults = { # format -> (codec, ext)
+_format_defaults = {  # format -> (codec, ext)
     'wav': ('pcm_s16le', 'wav'),
     'mp3': ('libmp3lame', 'mp3'),
     'flac': ('flac', 'flac'),
@@ -35,11 +34,13 @@ _extract_audio_param_types = [
     ts.VideoType(nullable=False),
     ts.IntType(nullable=False),
     ts.StringType(nullable=False),
-    ts.StringType(nullable=False)
+    ts.StringType(nullable=True),
 ]
+
+
 @func.udf(return_type=ts.AudioType(nullable=True), param_types=_extract_audio_param_types)
 def extract_audio(
-        video_path: str, stream_idx: int = 0, format: str = 'wav', codec: Optional[str] = None
+    video_path: str, stream_idx: int = 0, format: str = 'wav', codec: Optional[str] = None
 ) -> Optional[str]:
     """Extract an audio stream from a video file, save it as a media file and return its path"""
     if format not in _format_defaults:
@@ -51,12 +52,49 @@ def extract_audio(
             return None
         audio_stream = container.streams.audio[stream_idx]
         # create this in our tmp directory, so it'll get cleaned up if it's being generated as part of a query
-        output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+        output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
 
-        with av.open(output_filename, "w", format=format) as output_container:
+        with av.open(output_filename, 'w', format=format) as output_container:
             output_stream = output_container.add_stream(codec or default_codec)
             for packet in container.demux(audio_stream):
                 for frame in packet.decode():
                     output_container.mux(output_stream.encode(frame))
 
         return output_filename
+
+
+@func.udf(return_type=ts.JsonType(nullable=False), param_types=[ts.VideoType(nullable=False)])
+def get_metadata(video: str) -> dict:
+    """Gets various metadata associated with a video file.
+
+    Args:
+        video (str): Path to the video file.
+
+    Returns:
+        A dictionary containing the associated metadata.
+    """
+    with av.open(video) as container:
+        assert isinstance(container, av.container.InputContainer)
+        video_streams_info = [
+            {
+                'duration': stream.duration,
+                'frames': stream.frames,
+                'language': stream.language,
+                'average_rate': float(stream.average_rate) if stream.average_rate is not None else None,
+                'base_rate': float(stream.base_rate) if stream.base_rate is not None else None,
+                'guessed_rate': float(stream.guessed_rate) if stream.guessed_rate is not None else None,
+                'pix_fmt': getattr(stream.codec_context, 'pix_fmt', None),
+                'width': stream.width,
+                'height': stream.height,
+            }
+            for stream in container.streams
+            if isinstance(stream, av.video.stream.VideoStream)
+        ]
+        result = {
+            'bit_exact': container.bit_exact,
+            'bit_rate': container.bit_rate,
+            'size': container.size,
+            'metadata': container.metadata,
+            'streams': video_streams_info,  # TODO: Audio streams?
+        }
+    return result