<commit_before>#include <Python.h>
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/arrayobject.h>

#include <cstdint>
#include <memory>
#include <set>
#include <string>
#include <sstream>
#include <vector>

#include "caffe2/core/context.h"
#ifndef PYCAFFE2_CPU_ONLY
#include "caffe2/core/context_gpu.h"
#endif  // PYCAFFE2_CPU_ONLY
#include "caffe2/core/init.h"
#include "caffe2/core/net.h"
#include "caffe2/core/operator.h"
#include "caffe2/core/workspace.h"
#include "caffe2/proto/caffe2.pb.h"

//using namespace caffe2;  // NOLINT
using caffe2::Blob;
using caffe2::DeviceOption;
using caffe2::Tensor;
using caffe2::Workspace;
using caffe2::CPUContext;
using caffe2::OperatorDef;

#ifndef PYCAFFE2_CPU_ONLY
using caffe2::CUDAContext;
#endif  // PYCAFFE2_CPU_ONLY

// gWorkspaces allows us to define and switch between multiple workspaces in
// Python.
static std::map<std::string, std::unique_ptr<Workspace> > gWorkspaces;
// gWorkspace is the pointer to the current workspace. The ownership is kept
// by the gWorkspaces map.
static Workspace* gWorkspace = nullptr;
static std::string gCurrentWorkspaceName;

namespace {

using caffe2::string;

bool SwitchWorkspaceInternal(const string& name, const bool create_if_missing) {
  if (gWorkspaces.count(name)) {
    gCurrentWorkspaceName = name;
    gWorkspace = gWorkspaces[name].get();
    return true;
  } else if (create_if_missing) {
    std::unique_ptr<Workspace> new_workspace(new Workspace());
    gWorkspace = new_workspace.get();
    gWorkspaces.insert(std::make_pair(name, std::move(new_workspace)));
    gCurrentWorkspaceName = name;
    return true;
  } else {
    return false;
  }
}

inline string PyStringToStdString(PyObject* pystring) {
  return string(PyString_AsString(pystring), PyString_Size(pystring));
}

inline PyObject* StdStringToPyString(const string& str) {
  return PyString_FromStringAndSize(str.c_str(), str.size());
}

template <typename T>
inline void MakeStringInternal(std::stringstream& ss, const T& t) {
  ss << t;
}

template <typename T, typename ... Args>
inline void MakeStringInternal(std::stringstream& ss, const T& t, const Args&... args) {
  MakeStringInternal(ss, t);
  MakeStringInternal(ss, args...);
}

template <typename... Args>
string MakeString(const Args&... args) {
  std::stringstream ss;
  MakeStringInternal(ss, args...);
  return string(ss.str());
}


inline void PyErr_SetString(PyObject* type, const string& str) {
  PyErr_SetString(type, str.c_str());
}

static_assert(sizeof(int) == sizeof(int32_t),
              "Yangqing made a loose assumption that int will always be int32 "
              "for numpy type mapping");

template <typename T> struct NumpyTypeWrapper;
template<> struct NumpyTypeWrapper<float> {
  static const int type = NPY_FLOAT;
};
template<> struct NumpyTypeWrapper<int> {
  static const int type = NPY_INT32;
};

template <typename T, class DeviceContext>
PyObject* FetchTensor(const Tensor<DeviceContext>& tensor) {
  DeviceContext context;
  CAFFE_CHECK_GT(tensor.size(), 0);
  std::vector<npy_intp> npy_dims;
  for (const int dim : tensor.dims()) {
    npy_dims.push_back(dim);
  }
  PyObject* array = PyArray_SimpleNew(
      tensor.ndim(), npy_dims.data(), NumpyTypeWrapper<T>::type);
  // Now, copy the data to the tensor.
  // TODO(Yangqing): Is there an easier way to convert PyObject to
  // PyArrayObject?
  context.template Copy<T, DeviceContext, caffe2::CPUContext>(
      tensor.size(), tensor.template data<T>(),
      static_cast<T*>(PyArray_DATA(reinterpret_cast<PyArrayObject*>(array))));
  return array;
}

template <typename T, class DeviceContext>
PyObject* FeedTensor(const DeviceOption& option, PyArrayObject* original_array,
                     Blob* blob) {
  PyArrayObject* array = PyArray_GETCONTIGUOUS(original_array);
  DeviceContext context(option);
  Tensor<DeviceContext>* tensor =
      blob->GetMutable<Tensor<DeviceContext> >();
  // numpy requires long int as its dims.
  int ndim = PyArray_NDIM(array);
  npy_intp* npy_dims = PyArray_DIMS(array);
  std::vector<int> dims;
  for (int i = 0; i < ndim; ++i) {
    dims.push_back(npy_dims[i]);
  }
  tensor->Reshape(dims);
  // Now, copy the data to the tensor.
  context.template Copy<T, caffe2::CPUContext, DeviceContext>(
      tensor->size(), static_cast<T*>(PyArray_DATA(array)),
      tensor->template mutable_data<T>());
  Py_XDECREF(array);
  Py_RETURN_TRUE;
}

}  // namespace
<commit_msg>race condition fix: since Memcpy is now async, we will make sure that the python interface syncs before returning the content. Otherwise it makes things flaky.<commit_after>#include <Python.h>
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/arrayobject.h>

#include <cstdint>
#include <memory>
#include <set>
#include <string>
#include <sstream>
#include <vector>

#include "caffe2/core/context.h"
#ifndef PYCAFFE2_CPU_ONLY
#include "caffe2/core/context_gpu.h"
#endif  // PYCAFFE2_CPU_ONLY
#include "caffe2/core/init.h"
#include "caffe2/core/net.h"
#include "caffe2/core/operator.h"
#include "caffe2/core/workspace.h"
#include "caffe2/proto/caffe2.pb.h"

//using namespace caffe2;  // NOLINT
using caffe2::Blob;
using caffe2::DeviceOption;
using caffe2::Tensor;
using caffe2::Workspace;
using caffe2::CPUContext;
using caffe2::OperatorDef;

#ifndef PYCAFFE2_CPU_ONLY
using caffe2::CUDAContext;
#endif  // PYCAFFE2_CPU_ONLY

// gWorkspaces allows us to define and switch between multiple workspaces in
// Python.
static std::map<std::string, std::unique_ptr<Workspace> > gWorkspaces;
// gWorkspace is the pointer to the current workspace. The ownership is kept
// by the gWorkspaces map.
static Workspace* gWorkspace = nullptr;
static std::string gCurrentWorkspaceName;

namespace {

using caffe2::string;

bool SwitchWorkspaceInternal(const string& name, const bool create_if_missing) {
  if (gWorkspaces.count(name)) {
    gCurrentWorkspaceName = name;
    gWorkspace = gWorkspaces[name].get();
    return true;
  } else if (create_if_missing) {
    std::unique_ptr<Workspace> new_workspace(new Workspace());
    gWorkspace = new_workspace.get();
    gWorkspaces.insert(std::make_pair(name, std::move(new_workspace)));
    gCurrentWorkspaceName = name;
    return true;
  } else {
    return false;
  }
}

inline string PyStringToStdString(PyObject* pystring) {
  return string(PyString_AsString(pystring), PyString_Size(pystring));
}

inline PyObject* StdStringToPyString(const string& str) {
  return PyString_FromStringAndSize(str.c_str(), str.size());
}

template <typename T>
inline void MakeStringInternal(std::stringstream& ss, const T& t) {
  ss << t;
}

template <typename T, typename ... Args>
inline void MakeStringInternal(std::stringstream& ss, const T& t, const Args&... args) {
  MakeStringInternal(ss, t);
  MakeStringInternal(ss, args...);
}

template <typename... Args>
string MakeString(const Args&... args) {
  std::stringstream ss;
  MakeStringInternal(ss, args...);
  return string(ss.str());
}


inline void PyErr_SetString(PyObject* type, const string& str) {
  PyErr_SetString(type, str.c_str());
}

static_assert(sizeof(int) == sizeof(int32_t),
              "Yangqing made a loose assumption that int will always be int32 "
              "for numpy type mapping");

template <typename T> struct NumpyTypeWrapper;
template<> struct NumpyTypeWrapper<float> {
  static const int type = NPY_FLOAT;
};
template<> struct NumpyTypeWrapper<int> {
  static const int type = NPY_INT32;
};

template <typename T, class DeviceContext>
PyObject* FetchTensor(const Tensor<DeviceContext>& tensor) {
  DeviceContext context;
  CAFFE_CHECK_GT(tensor.size(), 0);
  std::vector<npy_intp> npy_dims;
  for (const int dim : tensor.dims()) {
    npy_dims.push_back(dim);
  }
  PyObject* array = PyArray_SimpleNew(
      tensor.ndim(), npy_dims.data(), NumpyTypeWrapper<T>::type);
  // Now, copy the data to the tensor.
  // TODO(Yangqing): Is there an easier way to convert PyObject to
  // PyArrayObject?
  context.template Copy<T, DeviceContext, caffe2::CPUContext>(
      tensor.size(), tensor.template data<T>(),
      static_cast<T*>(PyArray_DATA(reinterpret_cast<PyArrayObject*>(array))));
  context.FinishDeviceComputation();
  return array;
}

template <typename T, class DeviceContext>
PyObject* FeedTensor(const DeviceOption& option, PyArrayObject* original_array,
                     Blob* blob) {
  PyArrayObject* array = PyArray_GETCONTIGUOUS(original_array);
  DeviceContext context(option);
  Tensor<DeviceContext>* tensor =
      blob->GetMutable<Tensor<DeviceContext> >();
  // numpy requires long int as its dims.
  int ndim = PyArray_NDIM(array);
  npy_intp* npy_dims = PyArray_DIMS(array);
  std::vector<int> dims;
  for (int i = 0; i < ndim; ++i) {
    dims.push_back(npy_dims[i]);
  }
  tensor->Reshape(dims);
  // Now, copy the data to the tensor.
  context.template Copy<T, caffe2::CPUContext, DeviceContext>(
      tensor->size(), static_cast<T*>(PyArray_DATA(array)),
      tensor->template mutable_data<T>());
  context.FinishDeviceComputation();
  Py_XDECREF(array);
  Py_RETURN_TRUE;
}

}  // namespace

{'fraction_non_alphanumeric': 0.1251649802023757, 'fraction_numerical': 0.007259128904531456, 'mean_word_length': 4.866451612903226, 'pattern_counts': {'":': 0, '<': 79, '<?xml version=': 0, '>': 81, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5492795', 'n_tokens_mistral': 3069, 'n_tokens_neox': 3051, 'n_words': 741}