pixeltable 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pixeltable/func/udf.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import inspect
+from typing import List, Callable, Optional, overload, Any
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from .batched_function import ExplicitBatchedFunction
+from .callable_function import CallableFunction
+from .expr_template_function import ExprTemplateFunction
+from .function import Function
+from .function_registry import FunctionRegistry
+from .signature import Signature
+
+
+# Decorator invoked without parentheses: @pxt.udf
+@overload
+def udf(decorated_fn: Callable) -> Function: ...
+
+
+# Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
+@overload
+def udf(
+        *,
+        return_type: Optional[ts.ColumnType] = None,
+        param_types: Optional[List[ts.ColumnType]] = None,
+        batch_size: Optional[int] = None,
+        substitute_fn: Optional[Callable] = None,
+        _force_stored: bool = False
+) -> Callable: ...
+
+
+def udf(*args, **kwargs):
+    """A decorator to create a Function from a function definition.
+
+    Examples:
+        >>> @pxt.udf
+        ... def my_function(x: int) -> int:
+        ...    return x + 1
+
+        >>> @pxt.udf(param_types=[pxt.IntType()], return_type=pxt.IntType())
+        ... def my_function(x):
+        ...    return x + 1
+    """
+    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+
+        # Decorator invoked without parentheses: @pxt.udf
+        # Simply call make_function with defaults.
+        return make_function(decorated_fn=args[0])
+
+    else:
+
+        # Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
+        # Create a decorator for the specified schema.
+        return_type = kwargs.pop('return_type', None)
+        param_types = kwargs.pop('param_types', None)
+        batch_size = kwargs.pop('batch_size', None)
+        substitute_fn = kwargs.pop('py_fn', None)
+        force_stored = kwargs.pop('_force_stored', False)
+
+        def decorator(decorated_fn: Callable):
+            return make_function(
+                decorated_fn, return_type, param_types, batch_size, substitute_fn=substitute_fn,
+                force_stored=force_stored)
+
+        return decorator
+
+
+def make_function(
+    decorated_fn: Callable,
+    return_type: Optional[ts.ColumnType] = None,
+    param_types: Optional[List[ts.ColumnType]] = None,
+    batch_size: Optional[int] = None,
+    substitute_fn: Optional[Callable] = None,
+    function_name: Optional[str] = None,
+    force_stored: bool = False
+) -> Function:
+    """
+    Constructs a `CallableFunction` or `BatchedFunction`, depending on the
+    supplied parameters. If `substitute_fn` is specified, then `decorated_fn`
+    will be used only for its signature, with execution delegated to
+    `substitute_fn`.
+    """
+    # Obtain function_path from decorated_fn when appropriate
+    if force_stored:
+        # force storing the function in the db
+        function_path = None
+    elif decorated_fn.__module__ != '__main__' and decorated_fn.__name__.isidentifier():
+        function_path = f'{decorated_fn.__module__}.{decorated_fn.__qualname__}'
+    else:
+        function_path = None
+
+    # Derive function_name, if not specified explicitly
+    if function_name is None:
+        function_name = decorated_fn.__name__
+
+    # Display name to use for error messages
+    errmsg_name = function_name if function_path is None else function_path
+
+    sig = Signature.create(decorated_fn, param_types, return_type)
+
+    # batched functions must have a batched return type
+    # TODO: remove 'Python' from the error messages when we have full inference with Annotated types
+    if batch_size is not None and not sig.is_batched:
+        raise excs.Error(f'{errmsg_name}(): batch_size is specified; Python return type must be a `Batch`')
+    if batch_size is not None and len(sig.batched_parameters) == 0:
+        raise excs.Error(f'{errmsg_name}(): batch_size is specified; at least one Python parameter must be `Batch`')
+    if batch_size is None and len(sig.batched_parameters) > 0:
+        raise excs.Error(f'{errmsg_name}(): batched parameters in udf, but no `batch_size` given')
+
+    if substitute_fn is None:
+        py_fn = decorated_fn
+    else:
+        if function_path is None:
+            raise excs.Error(f'{errmsg_name}(): @udf decorator with a `substitute_fn` can only be used in a module')
+        py_fn = substitute_fn
+
+    if batch_size is None:
+        result = CallableFunction(signature=sig, py_fn=py_fn, self_path=function_path, self_name=function_name)
+    else:
+        result = ExplicitBatchedFunction(
+            signature=sig, batch_size=batch_size, invoker_fn=py_fn, self_path=function_path)
+
+    # If this function is part of a module, register it
+    if function_path is not None:
+        FunctionRegistry.get().register_function(function_path, result)
+
+    return result
+
+@overload
+def expr_udf(py_fn: Callable) -> ExprTemplateFunction: ...
+
+@overload
+def expr_udf(*, param_types: Optional[List[ts.ColumnType]] = None) -> Callable: ...
+
+def expr_udf(*args: Any, **kwargs: Any) -> Any:
+    def decorator(py_fn: Callable, param_types: Optional[List[ts.ColumnType]]) -> ExprTemplateFunction:
+        if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
+            # this is a named function in a module
+            function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
+        else:
+            function_path = None
+
+        sig = Signature.create(py_fn, param_types=param_types, return_type=None)
+        # TODO: verify that the inferred return type matches that of the template
+        # TODO: verify that the signature doesn't contain batched parameters
+
+        # construct Parameters from the function signature
+        import pixeltable.exprs as exprs
+        var_exprs = [exprs.Variable(param.name, param.col_type) for param in sig.parameters.values()]
+        # call the function with the parameter expressions to construct an Expr with parameters
+        template = py_fn(*var_exprs)
+        assert isinstance(template, exprs.Expr)
+        py_sig = inspect.signature(py_fn)
+        return ExprTemplateFunction(template, py_signature=py_sig, self_path=function_path, name=py_fn.__name__)
+
+    if len(args) == 1:
+        assert len(kwargs) == 0 and callable(args[0])
+        return decorator(args[0], None)
+    else:
+        assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
+        return lambda py_fn: decorator(py_fn, kwargs['param_types'])

pixeltable/functions/__init__.py

@@ -1,89 +1,57 @@
-import os
-from typing import Callable, List, Optional, Union
-import inspect
-from pathlib import Path
 import tempfile
+from pathlib import Path
+from typing import Optional, Union
 
-import PIL, cv2
+import PIL.Image
+import av
+import av.container
+import av.stream
 import numpy as np
 
-from pixeltable.type_system import StringType, IntType, JsonType, ColumnType, FloatType, ImageType, VideoType
-from pixeltable.function import Function
-from pixeltable import catalog
+import pixeltable.env as env
+import pixeltable.func as func
+# import all standard function modules here so they get registered with the FunctionRegistry
+import pixeltable.functions.pil.image
 from pixeltable import exprs
-from pixeltable import env
-import pixeltable.exceptions as exc
-
-
-def udf_call(eval_fn: Callable, return_type: ColumnType, tbl: Optional[catalog.Table]) -> exprs.FunctionCall:
-    """
-    Interprets eval_fn's parameters to be references to columns in 'tbl' and construct ColumnRefs as args.
-    """
-    params = inspect.signature(eval_fn).parameters
-    if len(params) > 0 and tbl is None:
-        raise exc.OperationalError(f'udf_call() is missing tbl parameter')
-    args: List[exprs.ColumnRef] = []
-    for param_name in params:
-        if param_name not in tbl.cols_by_name:
-            raise exc.OperationalError(
-                (f'udf_call(): lambda argument names need to be valid column names in table {tbl.name}: '
-                 f'column {param_name} unknown'))
-        args.append(exprs.ColumnRef(tbl.cols_by_name[param_name]))
-    fn = Function(return_type, [arg.col_type for arg in args], eval_fn=eval_fn)
-    return exprs.FunctionCall(fn, args)
+from pixeltable.type_system import IntType, ColumnType, FloatType, ImageType, VideoType
+# automatically import all submodules so that the udfs get registered
+from . import image, string, video, openai, together, fireworks, huggingface
 
+# TODO: remove and replace calls with astype()
 def cast(expr: exprs.Expr, target_type: ColumnType) -> exprs.Expr:
     expr.col_type = target_type
     return expr
 
-dict_map = Function(IntType(), [StringType(), JsonType()], eval_fn=lambda s, d: d[s])
-
-
-class SumAggregator:
+@func.uda(
+    update_types=[IntType()], value_type=IntType(), name='sum', allows_window=True, requires_order_by=False)
+class SumAggregator(func.Aggregator):
     def __init__(self):
         self.sum: Union[int, float] = 0
-    @classmethod
-    def make_aggregator(cls) -> 'SumAggregator':
-        return cls()
     def update(self, val: Union[int, float]) -> None:
         if val is not None:
             self.sum += val
     def value(self) -> Union[int, float]:
         return self.sum
 
-sum = Function(
-    IntType(), [IntType()],
-    module_name='pixeltable.functions',
-    init_symbol='SumAggregator.make_aggregator',
-    update_symbol='SumAggregator.update',
-    value_symbol='SumAggregator.value')
 
-class CountAggregator:
+@func.uda(
+    update_types=[IntType()], value_type=IntType(), name='count', allows_window = True, requires_order_by = False)
+class CountAggregator(func.Aggregator):
     def __init__(self):
         self.count = 0
-    @classmethod
-    def make_aggregator(cls) -> 'CountAggregator':
-        return cls()
     def update(self, val: int) -> None:
         if val is not None:
             self.count += 1
     def value(self) -> int:
         return self.count
 
-count = Function(
-    IntType(), [IntType()],
-    module_name = 'pixeltable.functions',
-    init_symbol = 'CountAggregator.make_aggregator',
-    update_symbol = 'CountAggregator.update',
-    value_symbol = 'CountAggregator.value')
 
-class MeanAggregator:
+@func.uda(
+    update_types=[IntType()], value_type=FloatType(), name='mean', allows_window=False, requires_order_by=False)
+class MeanAggregator(func.Aggregator):
     def __init__(self):
         self.sum = 0
         self.count = 0
-    @classmethod
-    def make_aggregator(cls) -> 'MeanAggregator':
-        return cls()
     def update(self, val: int) -> None:
         if val is not None:
             self.sum += val
@@ -93,54 +61,35 @@ class MeanAggregator:
             return None
         return self.sum / self.count
 
-mean = Function(
-    FloatType(), [IntType()],
-    module_name = 'pixeltable.functions',
-    init_symbol = 'MeanAggregator.make_aggregator',
-    update_symbol = 'MeanAggregator.update',
-    value_symbol = 'MeanAggregator.value')
 
-class VideoAggregator:
-    def __init__(self):
-        self.video_writer = None
-        self.size = None
-
-    @classmethod
-    def make_aggregator(cls) -> 'VideoAggregator':
-        return cls()
-
-    def update(self, frame_idx: int, frame: PIL.Image.Image) -> None:
-        if self.video_writer is None:
-            self.size = (frame.width, frame.height)
-            self.out_file = Path(os.getcwd()) / f'{Path(tempfile.mktemp()).name}.mp4'
-            self.tmp_file = Path(os.getcwd()) / f'{Path(tempfile.mktemp()).name}.mp4'
-            self.video_writer = cv2.VideoWriter(str(self.tmp_file), cv2.VideoWriter_fourcc(*'MP4V'), 25, self.size)
-
-        frame_array = np.array(frame)
-        frame_array = cv2.cvtColor(frame_array, cv2.COLOR_RGB2BGR)
-        self.video_writer.write(frame_array)
+@func.uda(
+    init_types=[IntType()], update_types=[ImageType()], value_type=VideoType(), name='make_video',
+    requires_order_by=True, allows_window=False)
+class VideoAggregator(func.Aggregator):
+    def __init__(self, fps: int = 25):
+        """follows https://pyav.org/docs/develop/cookbook/numpy.html#generating-video"""
+        self.container: Optional[av.container.OutputContainer] = None
+        self.stream: Optional[av.stream.Stream] = None
+        self.fps = fps
+
+    def update(self, frame: PIL.Image.Image) -> None:
+        if frame is None:
+            return
+        if self.container is None:
+            (_, output_filename) = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
+            self.out_file = Path(output_filename)
+            self.container = av.open(str(self.out_file), mode='w')
+            self.stream = self.container.add_stream('h264', rate=self.fps)
+            self.stream.pix_fmt = 'yuv420p'
+            self.stream.width = frame.width
+            self.stream.height = frame.height
+
+        av_frame = av.VideoFrame.from_ndarray(np.array(frame.convert('RGB')), format='rgb24')
+        for packet in self.stream.encode(av_frame):
+            self.container.mux(packet)
 
     def value(self) -> str:
-        self.video_writer.release()
-        os.system(f'ffmpeg -i {self.tmp_file} -vcodec libx264 {self.out_file}')
-        os.remove(self.tmp_file)
-        return self.out_file
-
-make_video = Function(
-    VideoType(), [IntType(), ImageType()],  # params: frame_idx, frame
-    order_by=[0],  # update() wants frames in frame_idx order
-    module_name = 'pixeltable.functions',
-    init_symbol = 'VideoAggregator.make_aggregator',
-    update_symbol = 'VideoAggregator.update',
-    value_symbol = 'VideoAggregator.value')
-
-
-__all__ = [
-    udf_call,
-    cast,
-    dict_map,
-    sum,
-    count,
-    mean,
-    make_video
-]
+        for packet in self.stream.encode():
+            self.container.mux(packet)
+        self.container.close()
+        return str(self.out_file)

pixeltable/functions/eval.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+from typing import List, Tuple, Dict
+from collections import defaultdict
+import sys
+
+import numpy as np
+
+import pixeltable.type_system as ts
+import pixeltable.func as func
+
+
+# TODO: figure out a better submodule structure
+
+# the following function has been adapted from MMEval
+# (sources at https://github.com/open-mmlab/mmeval)
+# Copyright (c) OpenMMLab. All rights reserved.
+def calculate_bboxes_area(bboxes: np.ndarray) -> np.ndarray:
+    """Calculate area of bounding boxes.
+
+    Args:
+        bboxes (numpy.ndarray): The bboxes with shape (n, 4) or (4, ) in 'xyxy' format.
+     Returns:
+        numpy.ndarray: The area of bboxes.
+    """
+    bboxes_w = (bboxes[..., 2] - bboxes[..., 0])
+    bboxes_h = (bboxes[..., 3] - bboxes[..., 1])
+    areas = bboxes_w * bboxes_h
+    return areas
+
+# the following function has been adapted from MMEval
+# (sources at https://github.com/open-mmlab/mmeval)
+# Copyright (c) OpenMMLab. All rights reserved.
+def calculate_overlaps(bboxes1: np.ndarray, bboxes2: np.ndarray) -> np.ndarray:
+    """Calculate the overlap between each bbox of bboxes1 and bboxes2.
+
+    Args:
+        bboxes1 (numpy.ndarray): The bboxes with shape (n, 4) in 'xyxy' format.
+        bboxes2 (numpy.ndarray): The bboxes with shape (k, 4) in 'xyxy' format.
+    Returns:
+        numpy.ndarray: IoUs or IoFs with shape (n, k).
+    """
+    bboxes1 = bboxes1.astype(np.float32)
+    bboxes2 = bboxes2.astype(np.float32)
+    rows = bboxes1.shape[0]
+    cols = bboxes2.shape[0]
+    overlaps = np.zeros((rows, cols), dtype=np.float32)
+
+    if rows * cols == 0:
+        return overlaps
+
+    if bboxes1.shape[0] > bboxes2.shape[0]:
+        # Swap bboxes for faster calculation.
+        bboxes1, bboxes2 = bboxes2, bboxes1
+        overlaps = np.zeros((cols, rows), dtype=np.float32)
+        exchange = True
+    else:
+        exchange = False
+
+    # Calculate the bboxes area.
+    area1 = calculate_bboxes_area(bboxes1)
+    area2 = calculate_bboxes_area(bboxes2)
+    eps = np.finfo(np.float32).eps
+
+    for i in range(bboxes1.shape[0]):
+        x_start = np.maximum(bboxes1[i, 0], bboxes2[:, 0])
+        y_start = np.maximum(bboxes1[i, 1], bboxes2[:, 1])
+        x_end = np.minimum(bboxes1[i, 2], bboxes2[:, 2])
+        y_end = np.minimum(bboxes1[i, 3], bboxes2[:, 3])
+        overlap_w = np.maximum(x_end - x_start, 0)
+        overlap_h = np.maximum(y_end - y_start, 0)
+        overlap = overlap_w * overlap_h
+
+        union = area1[i] + area2 - overlap
+        union = np.maximum(union, eps)
+        overlaps[i, :] = overlap / union
+    return overlaps if not exchange else overlaps.T
+
+
+# the following function has been adapted from MMEval
+# (sources at https://github.com/open-mmlab/mmeval)
+# Copyright (c) OpenMMLab. All rights reserved.
+def calculate_image_tpfp(
+    pred_bboxes: np.ndarray, pred_scores: np.ndarray, gt_bboxes: np.ndarray, min_iou: float
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Calculate the true positive and false positive on an image.
+
+    Args:
+        pred_bboxes (numpy.ndarray): Predicted bboxes of this image, with
+            shape (N, 5). The scores The predicted score of the bbox is
+            concatenated behind the predicted bbox.
+        gt_bboxes (numpy.ndarray): Ground truth bboxes of this image, with
+            shape (M, 4).
+        min_iou (float): The IoU threshold.
+
+    Returns:
+        tuple (tp, fp):
+
+        - tp (numpy.ndarray): Shape (N,),
+          the true positive flag of each predicted bbox on this image.
+        - fp (numpy.ndarray): Shape (N,),
+          the false positive flag of each predicted bbox on this image.
+    """
+    # Step 1. Concatenate `gt_bboxes` and `ignore_gt_bboxes`, then set
+    # the `ignore_gt_flags`.
+    # all_gt_bboxes = np.concatenate((gt_bboxes, ignore_gt_bboxes))
+    # ignore_gt_flags = np.concatenate((np.zeros(
+    #     (gt_bboxes.shape[0], 1),
+    #     dtype=bool), np.ones((ignore_gt_bboxes.shape[0], 1), dtype=bool)))
+
+    # Step 2. Initialize the `tp` and `fp` arrays.
+    num_preds = pred_bboxes.shape[0]
+    tp = np.zeros(num_preds, dtype=np.int8)
+    fp = np.zeros(num_preds, dtype=np.int8)
+
+    # Step 3. If there are no gt bboxes in this image, then all pred bboxes
+    # within area range are false positives.
+    if gt_bboxes.shape[0] == 0:
+        fp[...] = 1
+        return tp, fp
+
+    # Step 4. Calculate the IoUs between the predicted bboxes and the
+    # ground truth bboxes.
+    ious = calculate_overlaps(pred_bboxes, gt_bboxes)
+    # For each pred bbox, the max iou with all gts.
+    ious_max = ious.max(axis=1)
+    # For each pred bbox, which gt overlaps most with it.
+    ious_argmax = ious.argmax(axis=1)
+    # Sort all pred bbox in descending order by scores.
+    sorted_indices = np.argsort(-pred_scores)
+
+    # Step 5. Count the `tp` and `fp` of each iou threshold and area range.
+    # The flags that gt bboxes have been matched.
+    gt_covered_flags = np.zeros(gt_bboxes.shape[0], dtype=bool)
+
+    # Count the prediction bboxes in order of decreasing score.
+    for pred_bbox_idx in sorted_indices:
+        if ious_max[pred_bbox_idx] >= min_iou:
+            matched_gt_idx = ious_argmax[pred_bbox_idx]
+            if not gt_covered_flags[matched_gt_idx]:
+                tp[pred_bbox_idx] = 1
+                gt_covered_flags[matched_gt_idx] = True
+            else:
+                # This gt bbox has been matched and counted as fp.
+                fp[pred_bbox_idx] = 1
+        else:
+            fp[pred_bbox_idx] = 1
+
+    return tp, fp
+
+@func.udf(
+    return_type=ts.JsonType(nullable=False),
+    param_types=[
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False)
+    ])
+def eval_detections(
+        pred_bboxes: List[List[int]], pred_classes: List[int], pred_scores: List[float],
+        gt_bboxes: List[List[int]], gt_classes: List[int]
+) -> Dict:
+    class_idxs = list(set(pred_classes + gt_classes))
+    result: List[Dict] = []
+    pred_bboxes_arr = np.asarray(pred_bboxes)
+    pred_classes_arr = np.asarray(pred_classes)
+    pred_scores_arr = np.asarray(pred_scores)
+    gt_bboxes_arr = np.asarray(gt_bboxes)
+    gt_classes_arr = np.asarray(gt_classes)
+    for class_idx in class_idxs:
+        pred_filter = pred_classes_arr == class_idx
+        gt_filter = gt_classes_arr == class_idx
+        class_pred_scores = pred_scores_arr[pred_filter]
+        tp, fp = calculate_image_tpfp(
+            pred_bboxes_arr[pred_filter], class_pred_scores, gt_bboxes_arr[gt_filter], [0.5])
+        ordered_class_pred_scores = -np.sort(-class_pred_scores)
+        result.append({
+            'min_iou': 0.5, 'class': class_idx, 'tp': tp.tolist(), 'fp': fp.tolist(),
+            'scores': ordered_class_pred_scores.tolist(), 'num_gts': gt_filter.sum().item(),
+        })
+    return result
+
+@func.uda(
+    update_types=[ts.JsonType()], value_type=ts.JsonType(), name='mean_ap', allows_std_agg=True, allows_window=False)
+class MeanAPAggregator:
+    def __init__(self):
+        self.class_tpfp: Dict[int, List[Dict]] = defaultdict(list)
+
+    def update(self, eval_dicts: List[Dict]) -> None:
+        for eval_dict in eval_dicts:
+            class_idx = eval_dict['class']
+            self.class_tpfp[class_idx].append(eval_dict)
+
+    def value(self) -> Dict:
+        eps = np.finfo(np.float32).eps
+        result: Dict[int, float] = {}
+        for class_idx, tpfp in self.class_tpfp.items():
+            a1 = [x['tp'] for x in tpfp]
+            tp = np.concatenate([x['tp'] for x in tpfp], axis=0)
+            fp = np.concatenate([x['fp'] for x in tpfp], axis=0)
+            num_gts = np.sum([x['num_gts'] for x in tpfp])
+            scores = np.concatenate([np.asarray(x['scores']) for x in tpfp])
+            sorted_idxs = np.argsort(-scores)
+            tp_cumsum = tp[sorted_idxs].cumsum()
+            fp_cumsum = fp[sorted_idxs].cumsum()
+            precision = tp_cumsum / np.maximum(tp_cumsum + fp_cumsum, eps)
+            recall = tp_cumsum / np.maximum(num_gts, eps)
+
+            mrec = np.hstack((0, recall, 1))
+            mpre = np.hstack((0, precision, 0))
+            for i in range(mpre.shape[0] - 1, 0, -1):
+                mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
+            ind = np.where(mrec[1:] != mrec[:-1])[0]
+            ap = np.sum((mrec[ind + 1] - mrec[ind]) * mpre[ind + 1])
+            result[class_idx] = ap.item()
+        return result

pixeltable/functions/fireworks.py

@@ -0,0 +1,61 @@
+import logging
+import os
+from typing import Optional
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable import env
+
+
+@pxt.udf
+def chat_completions(
+        prompt: str,
+        model: str,
+        *,
+        max_tokens: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    initialize()
+    kwargs = {
+        'max_tokens': max_tokens,
+        'repetition_penalty': repetition_penalty,
+        'top_k': top_k,
+        'top_p': top_p,
+        'temperature': temperature
+    }
+    kwargs_not_none = dict(filter(lambda x: x[1] is not None, kwargs.items()))
+    print(kwargs_not_none)
+    return fireworks.client.Completion.create(
+        model=model,
+        prompt_or_messages=prompt,
+        **kwargs_not_none
+    ).dict()
+
+
+def initialize():
+    global _is_fireworks_initialized
+    if _is_fireworks_initialized:
+        return
+
+    _logger.info('Initializing Fireworks client.')
+
+    config = pxt.env.Env.get().config
+
+    if 'fireworks' in config and 'api_key' in config['fireworks']:
+        api_key = config['fireworks']['api_key']
+    else:
+        api_key = os.environ.get('FIREWORKS_API_KEY')
+    if api_key is None or api_key == '':
+        raise excs.Error('Fireworks client not initialized (no API key configured).')
+
+    import fireworks.client
+
+    fireworks.client.api_key = api_key
+    _is_fireworks_initialized = True
+
+
+_logger = logging.getLogger('pixeltable')
+_is_fireworks_initialized = False