pixeltable 0.4.0rc2__py3-none-any.whl → 0.4.1__py3-none-any.whl

pixeltable/functions/huggingface.py

@@ -51,7 +51,7 @@ def sentence_transformer(
     """
     env.Env.get().require_package('sentence_transformers')
     device = resolve_torch_device('auto')
-    from sentence_transformers import SentenceTransformer  # type: ignore
+    from sentence_transformers import SentenceTransformer
 
     # specifying the device, moves the model to device (gpu:cuda/mps, cpu)
     model = _lookup_model(model_id, SentenceTransformer, device=device, pass_device_to_create=True)
@@ -170,7 +170,7 @@ def clip(text: Batch[str], *, model_id: str) -> Batch[pxt.Array[(None,), pxt.Flo
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
     import torch
-    from transformers import CLIPModel, CLIPProcessor  # type: ignore
+    from transformers import CLIPModel, CLIPProcessor
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
@@ -395,19 +395,21 @@ def speech2text_for_conditional_generation(audio: pxt.Audio, *, model_id: str, l
     device = resolve_torch_device('auto', allow_mps=False)  # Doesn't seem to work on 'mps'; use 'cpu' instead
     import torch
     import torchaudio  # type: ignore[import-untyped]
-    from transformers import Speech2TextForConditionalGeneration, Speech2TextProcessor
+    from transformers import Speech2TextForConditionalGeneration, Speech2TextProcessor, Speech2TextTokenizer
 
     model = _lookup_model(model_id, Speech2TextForConditionalGeneration.from_pretrained, device=device)
     processor = _lookup_processor(model_id, Speech2TextProcessor.from_pretrained)
+    tokenizer = processor.tokenizer
     assert isinstance(processor, Speech2TextProcessor)
+    assert isinstance(tokenizer, Speech2TextTokenizer)
 
-    if language is not None and language not in processor.tokenizer.lang_code_to_id:
+    if language is not None and language not in 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(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]
+    forced_bos_token_id: Optional[int] = None if language is None else tokenizer.lang_code_to_id[language]
 
     # Get the model's sampling rate. Default to 16 kHz (the standard) if not in config
     model_sampling_rate = getattr(model.config, 'sampling_rate', 16_000)

pixeltable/functions/llama_cpp.py

@@ -17,7 +17,7 @@ def create_chat_completion(
     model_path: Optional[str] = None,
     repo_id: Optional[str] = None,
     repo_filename: Optional[str] = None,
-    args: Optional[dict[str, Any]] = None,
+    model_kwargs: Optional[dict[str, Any]] = None,
 ) -> dict:
     """
     Generate a chat completion from a list of messages.
@@ -35,14 +35,14 @@ def create_chat_completion(
         repo_id: The Hugging Face model repo id (if using a pretrained model).
         repo_filename: A filename or glob pattern to match the model file in the repo (optional, if using a
             pretrained model).
-        args: Additional arguments to pass to the `create_chat_completions` call, such as `max_tokens`, `temperature`,
-            `top_p`, and `top_k`. For details, see the
+        model_kwargs: Additional keyword args for the llama_cpp `create_chat_completions` API, such as `max_tokens`,
+            `temperature`, `top_p`, and `top_k`. For details, see the
             [llama_cpp create_chat_completions documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
     """
     Env.get().require_package('llama_cpp', min_version=[0, 3, 1])
 
-    if args is None:
-        args = {}
+    if model_kwargs is None:
+        model_kwargs = {}
 
     if (model_path is None) == (repo_id is None):
         raise excs.Error('Exactly one of `model_path` or `repo_id` must be provided.')
@@ -56,7 +56,7 @@ def create_chat_completion(
     else:
         Env.get().require_package('huggingface_hub')
         llm = _lookup_pretrained_model(repo_id, repo_filename, n_gpu_layers)
-    return llm.create_chat_completion(messages, **args)  # type: ignore
+    return llm.create_chat_completion(messages, **model_kwargs)  # type: ignore
 
 
 def _is_gpu_available() -> bool:

pixeltable/functions/mistralai.py

@@ -5,7 +5,7 @@ first `pip install mistralai` and configure your Mistral AI credentials, as desc
 the [Working with Mistral AI](https://pixeltable.readme.io/docs/working-with-mistralai) tutorial.
 """
 
-from typing import TYPE_CHECKING, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 
@@ -16,7 +16,7 @@ from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import mistralai.types.basemodel
+    import mistralai
 
 
 @register_client('mistral')
@@ -32,16 +32,7 @@ def _mistralai_client() -> 'mistralai.Mistral':
 
 @pxt.udf(resource_pool='request-rate:mistral')
 async def chat_completions(
-    messages: list[dict[str, str]],
-    *,
-    model: str,
-    temperature: Optional[float] = 0.7,
-    top_p: Optional[float] = 1.0,
-    max_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    random_seed: Optional[int] = None,
-    response_format: Optional[dict] = None,
-    safe_prompt: Optional[bool] = False,
+    messages: list[dict[str, str]], *, model: str, model_kwargs: Optional[dict[str, Any]] = None
 ) -> dict:
     """
     Chat Completion API.
@@ -60,8 +51,8 @@ async def chat_completions(
     Args:
         messages: The prompt(s) to generate completions for.
         model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
-
-    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/chat>
+        model_kwargs: Additional keyword args for the Mistral `chat/completions` API.
+            For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/chat>
 
     Returns:
         A dictionary containing the response and other metadata.
@@ -73,34 +64,20 @@ async def chat_completions(
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl.add_computed_column(response=completions(messages, model='mistral-latest-small'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     Env.get().require_package('mistralai')
     result = await _mistralai_client().chat.complete_async(
         messages=messages,  # type: ignore[arg-type]
         model=model,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=_opt(max_tokens),
-        stop=stop,
-        random_seed=_opt(random_seed),
-        response_format=response_format,  # type: ignore[arg-type]
-        safe_prompt=safe_prompt,
+        **model_kwargs,
     )
     return result.dict()
 
 
 @pxt.udf(resource_pool='request-rate:mistral')
-async def fim_completions(
-    prompt: str,
-    *,
-    model: str,
-    temperature: Optional[float] = 0.7,
-    top_p: Optional[float] = 1.0,
-    max_tokens: Optional[int] = None,
-    min_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    random_seed: Optional[int] = None,
-    suffix: Optional[str] = None,
-) -> dict:
+async def fim_completions(prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
     """
     Fill-in-the-middle Completion API.
 
@@ -118,6 +95,8 @@ async def fim_completions(
     Args:
         prompt: The text/code to complete.
         model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
+        model_kwargs: Additional keyword args for the Mistral `fim/completions` API.
+            For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/fim>
 
     For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/fim>
 
@@ -130,18 +109,11 @@ async def fim_completions(
 
         >>> tbl.add_computed_column(response=completions(tbl.prompt, model='codestral-latest'))
     """
+    if model_kwargs is None:
+        model_kwargs = {}
+
     Env.get().require_package('mistralai')
-    result = await _mistralai_client().fim.complete_async(
-        prompt=prompt,
-        model=model,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=_opt(max_tokens),
-        min_tokens=_opt(min_tokens),
-        stop=stop,
-        random_seed=_opt(random_seed),
-        suffix=_opt(suffix),
-    )
+    result = await _mistralai_client().fim.complete_async(prompt=prompt, model=model, **model_kwargs)
     return result.dict()
 
 
@@ -182,15 +154,6 @@ def _(model: str) -> ts.ArrayType:
     return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
-_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
-
-
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/ollama.py

@@ -50,7 +50,7 @@ def generate(
         template: Prompt template to use.
         context: The context parameter returned from a previous call to `generate()`.
         raw: If `True`, no formatting will be applied to the prompt.
-        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and
+        options: Additional options for the Ollama `chat` call, such as `max_tokens`, `temperature`, `top_p`, and
             `top_k`. For details, see the
             [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
             section of the Ollama documentation.