pixeltable 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

pixeltable/functions/huggingface.py

@@ -46,7 +46,7 @@ def sentence_transformer(
         Add a computed column that applies the model `all-mpnet-base-2` to an existing Pixeltable column `tbl.sentence`
         of the table `tbl`:
 
-        >>> tbl['result'] = sentence_transformer(tbl.sentence, model_id='all-mpnet-base-v2')
+        >>> tbl.add_computed_column(result=sentence_transformer(tbl.sentence, model_id='all-mpnet-base-v2'))
     """
     env.Env.get().require_package('sentence_transformers')
     device = resolve_torch_device('auto')
@@ -111,9 +111,9 @@ def cross_encoder(sentences1: Batch[str], sentences2: Batch[str], *, model_id: s
         Add a computed column that applies the model `ms-marco-MiniLM-L-4-v2` to the sentences in
         columns `tbl.sentence1` and `tbl.sentence2`:
 
-        >>> tbl['result'] = sentence_transformer(
-                tbl.sentence1, tbl.sentence2, model_id='ms-marco-MiniLM-L-4-v2'
-            )
+        >>> tbl.add_computed_column(result=sentence_transformer(
+        ...     tbl.sentence1, tbl.sentence2, model_id='ms-marco-MiniLM-L-4-v2'
+        ... ))
     """
     env.Env.get().require_package('sentence_transformers')
     device = resolve_torch_device('auto')
@@ -215,11 +215,7 @@ def _(model_id: str) -> pxt.ArrayType:
 
 @pxt.udf(batch_size=4)
 def detr_for_object_detection(
-    image: Batch[PIL.Image.Image],
-    *,
-    model_id: str,
-    threshold: float = 0.5,
-    revision: str = 'no_timm',
+    image: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5, revision: str = 'no_timm'
 ) -> Batch[dict]:
     """
     Computes DETR object detections for the specified image. `model_id` should be a reference to a pretrained
@@ -250,11 +246,11 @@ def detr_for_object_detection(
         Add a computed column that applies the model `facebook/detr-resnet-50` to an existing
         Pixeltable column `image` of the table `tbl`:
 
-        >>> tbl['detections'] = detr_for_object_detection(
+        >>> tbl.add_computed_column(detections=detr_for_object_detection(
         ...     tbl.image,
         ...     model_id='facebook/detr-resnet-50',
         ...     threshold=0.8
-        ... )
+        ... ))
     """
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
@@ -287,10 +283,7 @@ def detr_for_object_detection(
 
 @pxt.udf(batch_size=4)
 def vit_for_image_classification(
-    image: Batch[PIL.Image.Image],
-    *,
-    model_id: str,
-    top_k: int = 5
+    image: Batch[PIL.Image.Image], *, model_id: str, top_k: int = 5
 ) -> Batch[dict[str, Any]]:
     """
     Computes image classifications for the specified image using a Vision Transformer (ViT) model.
@@ -326,11 +319,11 @@ def vit_for_image_classification(
         Add a computed column that applies the model `google/vit-base-patch16-224` to an existing
         Pixeltable column `image` of the table `tbl`, returning the 10 most likely classes for each image:
 
-        >>> tbl['image_class'] = vit_for_image_classification(
+        >>> tbl.add_computed_column(image_class=vit_for_image_classification(
         ...     tbl.image,
         ...     model_id='google/vit-base-patch16-224',
         ...     top_k=10
-        ... )
+        ... ))
     """
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
@@ -362,12 +355,7 @@ def vit_for_image_classification(
 
 
 @pxt.udf
-def speech2text_for_conditional_generation(
-    audio: pxt.Audio,
-    *,
-    model_id: str,
-    language: Optional[str] = None,
-) -> str:
+def speech2text_for_conditional_generation(audio: pxt.Audio, *, model_id: str, language: Optional[str] = None) -> str:
     """
     Transcribes or translates speech to text using a Speech2Text model. `model_id` should be a reference to a
     pretrained [Speech2Text](https://huggingface.co/docs/transformers/en/model_doc/speech_to_text) model.
@@ -390,19 +378,19 @@ def speech2text_for_conditional_generation(
         Add a computed column that applies the model `facebook/s2t-small-librispeech-asr` to an existing
         Pixeltable column `audio` of the table `tbl`:
 
-        >>> tbl['transcription'] = speech2text_for_conditional_generation(
+        >>> tbl.add_computed_column(transcription=speech2text_for_conditional_generation(
         ...     tbl.audio,
         ...     model_id='facebook/s2t-small-librispeech-asr'
-        ... )
+        ... ))
 
         Add a computed column that applies the model `facebook/s2t-medium-mustc-multilingual-st` to an existing
         Pixeltable column `audio` of the table `tbl`, translating the audio to French:
 
-        >>> tbl['translation'] = speech2text_for_conditional_generation(
+        >>> tbl.add_computed_column(translation=speech2text_for_conditional_generation(
         ...     tbl.audio,
         ...     model_id='facebook/s2t-medium-mustc-multilingual-st',
         ...     language='fr'
-        ... )
+        ... ))
     """
     env.Env.get().require_package('transformers')
     env.Env.get().require_package('torchaudio')
@@ -419,7 +407,8 @@ def speech2text_for_conditional_generation(
     if language is not None and language not in processor.tokenizer.lang_code_to_id:
         raise excs.Error(
             f"Language code '{language}' is not supported by the model '{model_id}'. "
-            f"Supported languages are: {list(processor.tokenizer.lang_code_to_id.keys())}")
+            f'Supported languages are: {list(processor.tokenizer.lang_code_to_id.keys())}'
+        )
 
     forced_bos_token_id: Optional[int] = None if language is None else processor.tokenizer.lang_code_to_id[language]
 
@@ -439,11 +428,7 @@ def speech2text_for_conditional_generation(
     assert waveform.dim() == 1
 
     with torch.no_grad():
-        inputs = processor(
-            waveform,
-            sampling_rate=model_sampling_rate,
-            return_tensors='pt'
-        )
+        inputs = processor(waveform, sampling_rate=model_sampling_rate, return_tensors='pt')
         generated_ids = model.generate(**inputs.to(device), forced_bos_token_id=forced_bos_token_id).to('cpu')
 
     transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)
@@ -466,7 +451,7 @@ def detr_to_coco(image: PIL.Image.Image, detr_info: dict[str, Any]) -> dict[str,
         Add a computed column that converts the output `tbl.detections` to COCO format, where `tbl.image`
         is the image for which detections were computed:
 
-        >>> tbl['detections_coco'] = detr_to_coco(tbl.image, tbl.detections)
+        >>> tbl.add_computed_column(detections_coco=detr_to_coco(tbl.image, tbl.detections))
     """
     bboxes, labels = detr_info['boxes'], detr_info['labels']
     annotations = [
@@ -480,10 +465,7 @@ T = TypeVar('T')
 
 
 def _lookup_model(
-    model_id: str,
-    create: Callable[..., T],
-    device: Optional[str] = None,
-    pass_device_to_create: bool = False
+    model_id: str, create: Callable[..., T], device: Optional[str] = None, pass_device_to_create: bool = False
 ) -> T:
     from torch import nn
 

pixeltable/functions/image.py

@@ -56,6 +56,7 @@ def blend(im1: PIL.Image.Image, im2: PIL.Image.Image, alpha: float) -> PIL.Image
     """
     pass
 
+
 @pxt.udf(substitute_fn=PIL.Image.composite, is_method=True)
 def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.Image) -> PIL.Image.Image:
     """
@@ -68,6 +69,7 @@ def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.
 
 # PIL.Image.Image methods
 
+
 # Image.convert()
 @pxt.udf(is_method=True)
 def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
@@ -108,7 +110,9 @@ def _(self: Expr, box: tuple[int, int, int, int]) -> pxt.ColumnType:
     input_type = self.col_type
     assert isinstance(input_type, pxt.ImageType)
     if (isinstance(box, list) or isinstance(box, tuple)) and len(box) == 4 and all(isinstance(x, int) for x in box):
-        return pxt.ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+        return pxt.ImageType(
+            size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable
+        )
     return pxt.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
 
 
@@ -339,7 +343,7 @@ def getprojection(self: PIL.Image.Image) -> tuple[int]:
 
 @pxt.udf(substitute_fn=PIL.Image.Image.histogram, is_method=True)
 def histogram(
-        self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None
+    self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None
 ) -> list[int]:
     """
     Return a histogram for the image.
@@ -375,7 +379,7 @@ def quantize(
         kmeans: The number of k-means clusters to use.
         palette: The palette to use.
         dither: The dithering method. See the [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.quantize) for a list of supported methods.
-     """
+    """
     pass
 
 

pixeltable/functions/json.py

@@ -21,6 +21,7 @@ class make_list(pxt.Aggregator):
     """
     Collects arguments into a list.
     """
+
     def __init__(self) -> None:
         self.output: list[Any] = []
 

pixeltable/functions/llama_cpp.py

@@ -87,10 +87,7 @@ def _lookup_pretrained_model(repo_id: str, filename: Optional[str], n_gpu_layers
     key = (repo_id, filename, n_gpu_layers)
     if key not in _model_cache:
         llm = llama_cpp.Llama.from_pretrained(
-            repo_id=repo_id,
-            filename=filename,
-            n_gpu_layers=n_gpu_layers,
-            verbose=False,
+            repo_id=repo_id, filename=filename, n_gpu_layers=n_gpu_layers, verbose=False
         )
         _model_cache[key] = llm
     return _model_cache[key]

pixeltable/functions/mistralai.py

@@ -21,6 +21,7 @@ if TYPE_CHECKING:
 @register_client('mistral')
 def _(api_key: str) -> 'mistralai.Mistral':
     import mistralai
+
     return mistralai.Mistral(api_key=api_key)
 
 
@@ -28,8 +29,8 @@ def _mistralai_client() -> 'mistralai.Mistral':
     return Env.get().get_client('mistral')
 
 
-@pxt.udf
-def chat_completions(
+@pxt.udf(resource_pool='request-rate:mistral')
+async def chat_completions(
     messages: list[dict[str, str]],
     *,
     model: str,
@@ -47,6 +48,10 @@ def chat_completions(
     Equivalent to the Mistral AI `chat/completions` API endpoint.
     For additional details, see: <https://docs.mistral.ai/api/#tag/chat>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `mistral`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install mistralai`
@@ -65,10 +70,10 @@ def chat_completions(
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl['response'] = completions(messages, model='mistral-latest-small')
+        ... tbl.add_computed_column(response=completions(messages, model='mistral-latest-small'))
     """
     Env.get().require_package('mistralai')
-    return _mistralai_client().chat.complete(
+    result = await _mistralai_client().chat.complete_async(
         messages=messages,  # type: ignore[arg-type]
         model=model,
         temperature=temperature,
@@ -78,11 +83,12 @@ def chat_completions(
         random_seed=_opt(random_seed),
         response_format=response_format,  # type: ignore[arg-type]
         safe_prompt=safe_prompt,
-    ).dict()
+    )
+    return result.dict()
 
 
-@pxt.udf
-def fim_completions(
+@pxt.udf(resource_pool='request-rate:mistral')
+async def fim_completions(
     prompt: str,
     *,
     model: str,
@@ -100,6 +106,10 @@ def fim_completions(
     Equivalent to the Mistral AI `fim/completions` API endpoint.
     For additional details, see: <https://docs.mistral.ai/api/#tag/fim>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `mistral`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install mistralai`
@@ -117,10 +127,10 @@ def fim_completions(
         Add a computed column that applies the model `codestral-latest`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
-        >>> tbl['response'] = completions(tbl.prompt, model='codestral-latest')
+        >>> tbl.add_computed_column(response=completions(tbl.prompt, model='codestral-latest'))
     """
     Env.get().require_package('mistralai')
-    return _mistralai_client().fim.complete(
+    result = await _mistralai_client().fim.complete_async(
         prompt=prompt,
         model=model,
         temperature=temperature,
@@ -129,23 +139,26 @@ def fim_completions(
         min_tokens=_opt(min_tokens),
         stop=stop,
         random_seed=_opt(random_seed),
-        suffix=_opt(suffix)
-    ).dict()
+        suffix=_opt(suffix),
+    )
+    return result.dict()
 
 
-_embedding_dimensions_cache: dict[str, int] = {
-    'mistral-embed': 1024
-}
+_embedding_dimensions_cache: dict[str, int] = {'mistral-embed': 1024}
 
 
-@pxt.udf(batch_size=16)
-def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
+@pxt.udf(batch_size=16, resource_pool='request-rate:mistral')
+async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Embeddings API.
 
     Equivalent to the Mistral AI `embeddings` API endpoint.
     For additional details, see: <https://docs.mistral.ai/api/#tag/embeddings>
 
+    Request throttling:
+    Applies the rate limit set in the config (section `mistral`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install mistralai`
@@ -158,10 +171,7 @@ def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt
         An array representing the application of the given embedding to `input`.
     """
     Env.get().require_package('mistralai')
-    result = _mistralai_client().embeddings.create(
-        inputs=input,
-        model=model,
-    )
+    result = _mistralai_client().embeddings.create(inputs=input, model=model)
     return [np.array(data.embedding, dtype=np.float64) for data in result.data]
 
 
@@ -176,6 +186,7 @@ _T = TypeVar('_T')
 
 def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
     from mistralai.types import UNSET
+
     return arg if arg is not None else UNSET
 
 

pixeltable/functions/ollama.py

@@ -14,6 +14,7 @@ if TYPE_CHECKING:
 @env.register_client('ollama')
 def _(host: str) -> 'ollama.Client':
     import ollama
+
     return ollama.Client(host=host)
 
 
@@ -97,22 +98,12 @@ def chat(
     import ollama
 
     client = _ollama_client() or ollama
-    return client.chat(
-        model=model,
-        messages=messages,
-        tools=tools,
-        format=format,
-        options=options,
-    ).dict()  # type: ignore[call-overload]
+    return client.chat(model=model, messages=messages, tools=tools, format=format, options=options).dict()  # type: ignore[call-overload]
 
 
 @pxt.udf(batch_size=16)
 def embed(
-    input: Batch[str],
-    *,
-    model: str,
-    truncate: bool = True,
-    options: Optional[dict] = None,
+    input: Batch[str], *, model: str, truncate: bool = True, options: Optional[dict] = None
 ) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Generate embeddings from a model.
@@ -131,12 +122,7 @@ def embed(
     import ollama
 
     client = _ollama_client() or ollama
-    results = client.embed(
-        model=model,
-        input=input,
-        truncate=truncate,
-        options=options,
-    ).dict()
+    results = client.embed(model=model, input=input, truncate=truncate, options=options).dict()
     return [np.array(data, dtype=np.float64) for data in results['embeddings']]