Method predict

lit_nlp/lib/caching.py:276–329

(self,
              inputs: Iterable[JsonDict],
              progress_indicator: Optional[ProgressIndicator] = lambda x: x,
              **kw)

    return outputs

  def predict(self,
              inputs: Iterable[JsonDict],
              progress_indicator: Optional[ProgressIndicator] = lambda x: x,
              **kw) -> list[JsonDict]:
    inputs_as_list = list(inputs)

    if self._strict_id_validation:
      self._validate_ids(inputs_as_list)

    # Try to get results from the cache.
    input_keys = [self.key_fn(d) for d in inputs_as_list]
    if (none_keys := [k for k in input_keys if k is None]):
      logging.warning(
          "Attmepting to retrieve %d (of %d) predictions from the cache where"
          " the cache key is None - this can be from a missing or empty example"
          " id. These will call model.predict() on this and subsequent calls.",
          len(none_keys),
          len(input_keys),
      )
    if self._cache.pred_lock_key(input_keys):
      with self._cache.get_pred_lock(input_keys):
        cached_results = self._get_results_from_cache(input_keys)
    else:
      cached_results = self._get_results_from_cache(input_keys)

    # Make a single list of everything that wasn't found in the cache,
    # to actually run the model on these inputs.
    miss_idxs = [i for i, v in enumerate(cached_results) if v is None]
    misses = [inputs_as_list[i] for i in miss_idxs]
    if misses:
      logging.info("%s: %d misses out of %d inputs", self._log_prefix,
                   len(miss_idxs), len(cached_results))
    else:
      # If all results were already cached, return them.
      return cached_results

    with self._cache.get_pred_lock(input_keys):
      model_preds = list(self.wrapped.predict(progress_indicator(misses)))
      logging.info("Received %d predictions from model", len(model_preds))

      if len(model_preds) != len(misses):
        raise ValueError(f"Received {len(model_preds)} predictions, which does "
                         f"not match {len(misses)}, the number of inputs.")

      # Merge results back into the output list.
      with self._cache.lock:
        for i, orig_idx in enumerate(miss_idxs):
          self._cache.put(model_preds[i], self.key_fn(inputs_as_list[orig_idx]))
          cached_results[orig_idx] = model_preds[i]

      # Remove the prediction lock from the cache as the request is complete
      self._cache.delete_pred_lock(input_keys)

    return cached_results

  def _get_results_from_cache(self, input_keys: list[CacheKey]):
    with self._cache.lock: