pixeltable 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

pixeltable/functions/llama_cpp.py

@@ -0,0 +1,106 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.env import Env
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import llama_cpp
+
+
+@pxt.udf
+def create_chat_completion(
+    messages: list[dict],
+    *,
+    model_path: Optional[str] = None,
+    repo_id: Optional[str] = None,
+    repo_filename: Optional[str] = None,
+    args: Optional[dict[str, Any]] = None,
+) -> dict:
+    """
+    Generate a chat completion from a list of messages.
+
+    The model can be specified either as a local path, or as a repo_id and repo_filename that reference a pretrained
+    model on the Hugging Face model hub. Exactly one of `model_path` or `repo_id` must be provided; if `model_path`
+    is provided, then an optional `repo_filename` can also be specified.
+
+    For additional 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).
+
+    Args:
+        messages: A list of messages to generate a response for.
+        model_path: Path to the model (if using a local model).
+        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
+            [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_path is None) == (repo_id is None):
+        raise excs.Error('Exactly one of `model_path` or `repo_id` must be provided.')
+    if (repo_id is None) and (repo_filename is not None):
+        raise excs.Error('`repo_filename` can only be provided along with `repo_id`.')
+
+    n_gpu_layers = -1 if _is_gpu_available() else 0  # 0 = CPU only, -1 = offload all layers to GPU
+
+    if model_path is not None:
+        llm = _lookup_local_model(model_path, n_gpu_layers)
+    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
+
+
+def _is_gpu_available() -> bool:
+    import llama_cpp
+
+    global _IS_GPU_AVAILABLE
+    if _IS_GPU_AVAILABLE is None:
+        llama_cpp_path = Path(llama_cpp.__file__).parent
+        lib = llama_cpp.llama_cpp.load_shared_library('llama', llama_cpp_path / 'lib')
+        _IS_GPU_AVAILABLE = bool(lib.llama_supports_gpu_offload())
+
+    return _IS_GPU_AVAILABLE
+
+
+def _lookup_local_model(model_path: str, n_gpu_layers: int) -> 'llama_cpp.Llama':
+    import llama_cpp
+
+    key = (model_path, None, n_gpu_layers)
+    if key not in _model_cache:
+        llm = llama_cpp.Llama(model_path, n_gpu_layers=n_gpu_layers)
+        _model_cache[key] = llm
+    return _model_cache[key]
+
+
+def _lookup_pretrained_model(repo_id: str, filename: Optional[str], n_gpu_layers: int) -> 'llama_cpp.Llama':
+    import llama_cpp
+
+    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
+        )
+        _model_cache[key] = llm
+    return _model_cache[key]
+
+
+_model_cache: dict[tuple[str, str, int], Any] = {}
+_IS_GPU_AVAILABLE: Optional[bool] = None
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/mistralai.py

@@ -140,8 +140,8 @@ _embedding_dimensions_cache: dict[str, int] = {
 }
 
 
-@pxt.udf(batch_size=16, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
-def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
+@pxt.udf(batch_size=16)
+def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), float]]:
     """
     Embeddings API.
 

pixeltable/functions/ollama.py

@@ -0,0 +1,147 @@
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable import env
+from pixeltable.func import Batch
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import ollama
+
+
+@env.register_client('ollama')
+def _(host: str) -> 'ollama.Client':
+    import ollama
+    return ollama.Client(host=host)
+
+
+def _ollama_client() -> Optional['ollama.Client']:
+    try:
+        return env.Env.get().get_client('ollama')
+    except Exception:
+        return None
+
+
+@pxt.udf
+def generate(
+    prompt: str,
+    *,
+    model: str,
+    suffix: str = '',
+    system: str = '',
+    template: str = '',
+    context: Optional[list[int]] = None,
+    raw: bool = False,
+    format: str = '',
+    options: Optional[dict] = None,
+) -> dict:
+    """
+    Generate a response for a given prompt with a provided model.
+
+    Args:
+        prompt: The prompt to generate a response for.
+        model: The model name.
+        suffix: The text after the model response.
+        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        system: System message.
+        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 `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.
+    """
+    env.Env.get().require_package('ollama')
+    import ollama
+
+    client = _ollama_client() or ollama
+    return client.generate(
+        model=model,
+        prompt=prompt,
+        suffix=suffix,
+        system=system,
+        template=template,
+        context=context,
+        raw=raw,
+        format=format,
+        options=options,
+    )  # type: ignore[call-overload]
+
+
+@pxt.udf
+def chat(
+    messages: list[dict],
+    *,
+    model: str,
+    tools: Optional[list[dict]] = None,
+    format: str = '',
+    options: Optional[dict] = None,
+) -> dict:
+    """
+    Generate the next message in a chat with a provided model.
+
+    Args:
+        messages: The messages of the chat.
+        model: The model name.
+        tools: Tools for the model to use.
+        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        options: Additional options to pass to the `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.
+    """
+    env.Env.get().require_package('ollama')
+    import ollama
+
+    client = _ollama_client() or ollama
+    return client.chat(
+        model=model,
+        messages=messages,
+        tools=tools,
+        format=format,
+        options=options,
+    )  # type: ignore[call-overload]
+
+
+@pxt.udf(batch_size=16)
+def embed(
+    input: Batch[str],
+    *,
+    model: str,
+    truncate: bool = True,
+    options: Optional[dict] = None,
+) -> Batch[pxt.Array[(None,), pxt.Float]]:
+    """
+    Generate embeddings from a model.
+
+    Args:
+        input: The input text to generate embeddings for.
+        model: The model name.
+        truncate: Truncates the end of each input to fit within context length.
+            Returns error if false and context length is exceeded.
+        options: Additional options to pass to the `embed` call.
+            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.
+    """
+    env.Env.get().require_package('ollama')
+    import ollama
+
+    client = _ollama_client() or ollama
+    results = client.embed(
+        model=model,
+        input=input,
+        truncate=truncate,
+        options=options,  # type: ignore[arg-type]
+    )
+    return [np.array(data, dtype=np.float64) for data in results['embeddings']]
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/openai.py

@@ -16,7 +16,6 @@ import PIL.Image
 import tenacity
 
 import pixeltable as pxt
-import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.utils.code import local_public_names
@@ -51,10 +50,10 @@ def _retry(fn: Callable) -> Callable:
 # Audio Endpoints
 
 
-@pxt.udf(return_type=ts.AudioType())
+@pxt.udf
 def speech(
     input: str, *, model: str, voice: str, response_format: Optional[str] = None, speed: Optional[float] = None
-) -> str:
+) -> pxt.Audio:
     """
     Generates audio from the input text.
 
@@ -91,17 +90,9 @@ def speech(
     return output_filename
 
 
-@pxt.udf(
-    param_types=[
-        ts.AudioType(),
-        ts.StringType(),
-        ts.StringType(nullable=True),
-        ts.StringType(nullable=True),
-        ts.FloatType(nullable=True),
-    ]
-)
+@pxt.udf
 def transcriptions(
-    audio: str,
+    audio: pxt.Audio,
     *,
     model: str,
     language: Optional[str] = None,
@@ -140,8 +131,14 @@ def transcriptions(
     return transcription.dict()
 
 
-@pxt.udf(param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True), ts.FloatType(nullable=True)])
-def translations(audio: str, *, model: str, prompt: Optional[str] = None, temperature: Optional[float] = None) -> dict:
+@pxt.udf
+def translations(
+    audio: pxt.Audio,
+    *,
+    model: str,
+    prompt: Optional[str] = None,
+    temperature: Optional[float] = None
+) -> dict:
     """
     Translates audio into English.
 
@@ -304,10 +301,10 @@ _embedding_dimensions_cache: dict[str, int] = {
 }
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
+@pxt.udf(batch_size=32)
 def embeddings(
     input: Batch[str], *, model: str, dimensions: Optional[int] = None, user: Optional[str] = None
-) -> Batch[np.ndarray]:
+) -> Batch[pxt.Array[(None,), float]]:
     """
     Creates an embedding vector representing the input text.
 
@@ -342,13 +339,13 @@ def embeddings(
 
 
 @embeddings.conditional_return_type
-def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
+def _(model: str, dimensions: Optional[int] = None) -> pxt.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)
+            return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
         dimensions = _embedding_dimensions_cache.get(model, None)
-    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
+    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType(), nullable=False)
 
 
 #####################################
@@ -408,17 +405,17 @@ def image_generations(
 
 
 @image_generations.conditional_return_type
-def _(size: Optional[str] = None) -> ts.ImageType:
+def _(size: Optional[str] = None) -> pxt.ImageType:
     if size is None:
-        return ts.ImageType(size=(1024, 1024))
+        return pxt.ImageType(size=(1024, 1024))
     x_pos = size.find('x')
     if x_pos == -1:
-        return ts.ImageType()
+        return pxt.ImageType()
     try:
         width, height = int(size[:x_pos]), int(size[x_pos + 1 :])
     except ValueError:
-        return ts.ImageType()
-    return ts.ImageType(size=(width, height))
+        return pxt.ImageType()
+    return pxt.ImageType(size=(width, height))
 
 
 #####################################

pixeltable/functions/replicate.py

@@ -0,0 +1,72 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various endpoints from the Replicate API. In order to use them, you must
+first `pip install replicate` and configure your Replicate credentials, as described in
+the [Working with Replicate](https://pixeltable.readme.io/docs/working-with-replicate) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import replicate  # type: ignore[import-untyped]
+
+
+@register_client('replicate')
+def _(api_token: str) -> 'replicate.Client':
+    import replicate
+    return replicate.Client(api_token=api_token)
+
+
+def _replicate_client() -> 'replicate.Client':
+    return Env.get().get_client('replicate')
+
+
+@pxt.udf
+def run(
+    input: dict[str, Any],
+    *,
+    ref: str,
+) -> dict[str, Any]:
+    """
+    Run a model on Replicate.
+
+    For additional details, see: <https://replicate.com/docs/topics/models/run-a-model>
+
+    __Requirements:__
+
+    - `pip install replicate`
+
+    Args:
+        input: The input parameters for the model.
+        ref: The name of the model to run.
+
+    Returns:
+        The output of the model.
+
+    Examples:
+        Add a computed column that applies the model `meta/meta-llama-3-8b-instruct`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> input = {'system_prompt': 'You are a helpful assistant.', 'prompt': tbl.prompt}
+        ... tbl['response'] = run(input, ref='meta/meta-llama-3-8b-instruct')
+
+        Add a computed column that uses the model `black-forest-labs/flux-schnell`
+        to generate images from an existing Pixeltable column `tbl.prompt`:
+
+        >>> input = {'prompt': tbl.prompt, 'go_fast': True, 'megapixels': '1'}
+        ... tbl['response'] = run(input, ref='black-forest-labs/flux-schnell')
+        ... tbl['image'] = tbl.response.output[0].astype(pxt.Image)
+    """
+    Env.get().require_package('replicate')
+    return _replicate_client().run(ref, input, use_file_output=False)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__