617 lines
24 KiB
Rust
617 lines
24 KiB
Rust
use std::collections::BTreeSet;
|
|
use std::fmt::{self, Debug, Display};
|
|
use std::fs;
|
|
|
|
use crate::all_config;
|
|
use crate::all_config::AllConfig;
|
|
use anyhow::{bail, Context};
|
|
use bpr_thrift::data::DataRecord;
|
|
use bpr_thrift::prediction_service::BatchPredictionRequest;
|
|
use bpr_thrift::tensor::GeneralTensor;
|
|
use log::debug;
|
|
use ndarray::Array2;
|
|
use once_cell::sync::OnceCell;
|
|
use ort::tensor::InputTensor;
|
|
use prometheus::{HistogramOpts, HistogramVec};
|
|
use segdense::mapper::{FeatureMapper, MapReader};
|
|
use segdense::segdense_transform_spec_home_recap_2022::{DensificationTransformSpec, Root};
|
|
use segdense::util;
|
|
use thrift::protocol::{TBinaryInputProtocol, TSerializable};
|
|
use thrift::transport::TBufferChannel;
|
|
|
|
pub fn log_feature_match(
|
|
dr: &DataRecord,
|
|
seg_dense_config: &DensificationTransformSpec,
|
|
dr_type: String,
|
|
) {
|
|
// Note the following algorithm matches features from config using linear search.
|
|
// Also the record source is MinDataRecord. This includes only binary and continous features for now.
|
|
|
|
for (feature_id, feature_value) in dr.continuous_features.as_ref().unwrap() {
|
|
debug!(
|
|
"{} - Continous Datarecord => Feature ID: {}, Feature value: {}",
|
|
dr_type, feature_id, feature_value
|
|
);
|
|
for input_feature in &seg_dense_config.cont.input_features {
|
|
if input_feature.feature_id == *feature_id {
|
|
debug!("Matching input feature: {:?}", input_feature)
|
|
}
|
|
}
|
|
}
|
|
|
|
for feature_id in dr.binary_features.as_ref().unwrap() {
|
|
debug!(
|
|
"{} - Binary Datarecord => Feature ID: {}",
|
|
dr_type, feature_id
|
|
);
|
|
for input_feature in &seg_dense_config.binary.input_features {
|
|
if input_feature.feature_id == *feature_id {
|
|
debug!("Found input feature: {:?}", input_feature)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn log_feature_matches(drs: &Vec<DataRecord>, seg_dense_config: &DensificationTransformSpec) {
|
|
for dr in drs {
|
|
log_feature_match(dr, seg_dense_config, String::from("individual"));
|
|
}
|
|
}
|
|
|
|
pub trait Converter: Send + Sync + Debug + 'static + Display {
|
|
fn convert(&self, input: Vec<Vec<u8>>) -> (Vec<InputTensor>, Vec<usize>);
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
#[allow(dead_code)]
|
|
pub struct BatchPredictionRequestToTorchTensorConverter {
|
|
all_config: AllConfig,
|
|
seg_dense_config: Root,
|
|
all_config_path: String,
|
|
seg_dense_config_path: String,
|
|
feature_mapper: FeatureMapper,
|
|
user_embedding_feature_id: i64,
|
|
user_eng_embedding_feature_id: i64,
|
|
author_embedding_feature_id: i64,
|
|
discrete_features_to_report: BTreeSet<i64>,
|
|
continuous_features_to_report: BTreeSet<i64>,
|
|
discrete_feature_metrics: &'static HistogramVec,
|
|
continuous_feature_metrics: &'static HistogramVec,
|
|
}
|
|
|
|
impl Display for BatchPredictionRequestToTorchTensorConverter {
|
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
|
write!(
|
|
f,
|
|
"all_config_path: {}, seg_dense_config_path:{}",
|
|
self.all_config_path, self.seg_dense_config_path
|
|
)
|
|
}
|
|
}
|
|
|
|
impl BatchPredictionRequestToTorchTensorConverter {
|
|
pub fn new(
|
|
model_dir: &str,
|
|
model_version: &str,
|
|
reporting_feature_ids: Vec<(i64, &str)>,
|
|
register_metric_fn: Option<impl Fn(&HistogramVec)>,
|
|
) -> anyhow::Result<BatchPredictionRequestToTorchTensorConverter> {
|
|
let all_config_path = format!("{}/{}/all_config.json", model_dir, model_version);
|
|
let seg_dense_config_path = format!(
|
|
"{}/{}/segdense_transform_spec_home_recap_2022.json",
|
|
model_dir, model_version
|
|
);
|
|
let seg_dense_config = util::load_config(&seg_dense_config_path)?;
|
|
let all_config = all_config::parse(
|
|
&fs::read_to_string(&all_config_path)
|
|
.with_context(|| "error loading all_config.json - ")?,
|
|
)?;
|
|
|
|
let feature_mapper = util::load_from_parsed_config(seg_dense_config.clone())?;
|
|
|
|
let user_embedding_feature_id = Self::get_feature_id(
|
|
&all_config
|
|
.train_data
|
|
.seg_dense_schema
|
|
.renamed_features
|
|
.user_embedding,
|
|
&seg_dense_config,
|
|
);
|
|
let user_eng_embedding_feature_id = Self::get_feature_id(
|
|
&all_config
|
|
.train_data
|
|
.seg_dense_schema
|
|
.renamed_features
|
|
.user_eng_embedding,
|
|
&seg_dense_config,
|
|
);
|
|
let author_embedding_feature_id = Self::get_feature_id(
|
|
&all_config
|
|
.train_data
|
|
.seg_dense_schema
|
|
.renamed_features
|
|
.author_embedding,
|
|
&seg_dense_config,
|
|
);
|
|
static METRICS: OnceCell<(HistogramVec, HistogramVec)> = OnceCell::new();
|
|
let (discrete_feature_metrics, continuous_feature_metrics) = METRICS.get_or_init(|| {
|
|
let discrete = HistogramVec::new(
|
|
HistogramOpts::new(":navi:feature_id:discrete", "Discrete Feature ID values")
|
|
.buckets(Vec::from(&[
|
|
0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0,
|
|
120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 250.0,
|
|
300.0, 500.0, 1000.0, 10000.0, 100000.0,
|
|
] as &'static [f64])),
|
|
&["feature_id"],
|
|
)
|
|
.expect("metric cannot be created");
|
|
let continuous = HistogramVec::new(
|
|
HistogramOpts::new(
|
|
":navi:feature_id:continuous",
|
|
"continuous Feature ID values",
|
|
)
|
|
.buckets(Vec::from(&[
|
|
0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0, 120.0,
|
|
130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 250.0, 300.0, 500.0,
|
|
1000.0, 10000.0, 100000.0,
|
|
] as &'static [f64])),
|
|
&["feature_id"],
|
|
)
|
|
.expect("metric cannot be created");
|
|
register_metric_fn.map(|r| {
|
|
r(&discrete);
|
|
r(&continuous);
|
|
});
|
|
(discrete, continuous)
|
|
});
|
|
|
|
let mut discrete_features_to_report = BTreeSet::new();
|
|
let mut continuous_features_to_report = BTreeSet::new();
|
|
|
|
for (feature_id, feature_type) in reporting_feature_ids.iter() {
|
|
match *feature_type {
|
|
"discrete" => discrete_features_to_report.insert(feature_id.clone()),
|
|
"continuous" => continuous_features_to_report.insert(feature_id.clone()),
|
|
_ => bail!(
|
|
"Invalid feature type {} for reporting metrics!",
|
|
feature_type
|
|
),
|
|
};
|
|
}
|
|
|
|
Ok(BatchPredictionRequestToTorchTensorConverter {
|
|
all_config,
|
|
seg_dense_config,
|
|
all_config_path,
|
|
seg_dense_config_path,
|
|
feature_mapper,
|
|
user_embedding_feature_id,
|
|
user_eng_embedding_feature_id,
|
|
author_embedding_feature_id,
|
|
discrete_features_to_report,
|
|
continuous_features_to_report,
|
|
discrete_feature_metrics,
|
|
continuous_feature_metrics,
|
|
})
|
|
}
|
|
|
|
fn get_feature_id(feature_name: &str, seg_dense_config: &Root) -> i64 {
|
|
// given a feature name, we get the complex feature type id
|
|
for feature in &seg_dense_config.complex_feature_type_transform_spec {
|
|
if feature.full_feature_name == feature_name {
|
|
return feature.feature_id;
|
|
}
|
|
}
|
|
-1
|
|
}
|
|
|
|
fn parse_batch_prediction_request(bytes: Vec<u8>) -> BatchPredictionRequest {
|
|
// parse batch prediction request into a struct from byte array repr.
|
|
let mut bc = TBufferChannel::with_capacity(bytes.len(), 0);
|
|
bc.set_readable_bytes(&bytes);
|
|
let mut protocol = TBinaryInputProtocol::new(bc, true);
|
|
BatchPredictionRequest::read_from_in_protocol(&mut protocol).unwrap()
|
|
}
|
|
|
|
fn get_embedding_tensors(
|
|
&self,
|
|
bprs: &[BatchPredictionRequest],
|
|
feature_id: i64,
|
|
batch_size: &[usize],
|
|
) -> Array2<f32> {
|
|
// given an embedding feature id, extract the float tensor array into tensors.
|
|
let cols: usize = 200;
|
|
let rows: usize = batch_size[batch_size.len() - 1];
|
|
let total_size = rows * cols;
|
|
|
|
let mut working_set = vec![0 as f32; total_size];
|
|
let mut bpr_start = 0;
|
|
for (bpr, &bpr_end) in bprs.iter().zip(batch_size) {
|
|
if bpr.common_features.is_some() {
|
|
if bpr.common_features.as_ref().unwrap().tensors.is_some() {
|
|
if bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.tensors
|
|
.as_ref()
|
|
.unwrap()
|
|
.contains_key(&feature_id)
|
|
{
|
|
let source_tensor = bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.tensors
|
|
.as_ref()
|
|
.unwrap()
|
|
.get(&feature_id)
|
|
.unwrap();
|
|
let tensor = match source_tensor {
|
|
GeneralTensor::FloatTensor(float_tensor) =>
|
|
//Tensor::of_slice(
|
|
{
|
|
float_tensor
|
|
.floats
|
|
.iter()
|
|
.map(|x| x.into_inner() as f32)
|
|
.collect::<Vec<_>>()
|
|
}
|
|
_ => vec![0 as f32; cols],
|
|
};
|
|
|
|
// since the tensor is found in common feature, add it in all batches
|
|
for row in bpr_start..bpr_end {
|
|
for col in 0..cols {
|
|
working_set[row * cols + col] = tensor[col];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// find the feature in individual feature list and add to corresponding batch.
|
|
for (index, datarecord) in bpr.individual_features_list.iter().enumerate() {
|
|
if datarecord.tensors.is_some()
|
|
&& datarecord
|
|
.tensors
|
|
.as_ref()
|
|
.unwrap()
|
|
.contains_key(&feature_id)
|
|
{
|
|
let source_tensor = datarecord
|
|
.tensors
|
|
.as_ref()
|
|
.unwrap()
|
|
.get(&feature_id)
|
|
.unwrap();
|
|
let tensor = match source_tensor {
|
|
GeneralTensor::FloatTensor(float_tensor) => float_tensor
|
|
.floats
|
|
.iter()
|
|
.map(|x| x.into_inner() as f32)
|
|
.collect::<Vec<_>>(),
|
|
_ => vec![0 as f32; cols],
|
|
};
|
|
for col in 0..cols {
|
|
working_set[(bpr_start + index) * cols + col] = tensor[col];
|
|
}
|
|
}
|
|
}
|
|
bpr_start = bpr_end;
|
|
}
|
|
Array2::<f32>::from_shape_vec([rows, cols], working_set).unwrap()
|
|
}
|
|
|
|
// Todo : Refactor, create a generic version with different type and field accessors
|
|
// Example paramterize and then instiantiate the following
|
|
// (FLOAT --> FLOAT, DataRecord.continuous_feature)
|
|
// (BOOL --> INT64, DataRecord.binary_feature)
|
|
// (INT64 --> INT64, DataRecord.discrete_feature)
|
|
fn get_continuous(&self, bprs: &[BatchPredictionRequest], batch_ends: &[usize]) -> InputTensor {
|
|
// These need to be part of model schema
|
|
let rows: usize = batch_ends[batch_ends.len() - 1];
|
|
let cols: usize = 5293;
|
|
let full_size: usize = rows * cols;
|
|
let default_val = f32::NAN;
|
|
|
|
let mut tensor = vec![default_val; full_size];
|
|
|
|
let mut bpr_start = 0;
|
|
for (bpr, &bpr_end) in bprs.iter().zip(batch_ends) {
|
|
// Common features
|
|
if bpr.common_features.is_some()
|
|
&& bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.continuous_features
|
|
.is_some()
|
|
{
|
|
let common_features = bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.continuous_features
|
|
.as_ref()
|
|
.unwrap();
|
|
|
|
for feature in common_features {
|
|
match self.feature_mapper.get(feature.0) {
|
|
Some(f_info) => {
|
|
let idx = f_info.index_within_tensor as usize;
|
|
if idx < cols {
|
|
// Set value in each row
|
|
for r in bpr_start..bpr_end {
|
|
let flat_index: usize = r * cols + idx;
|
|
tensor[flat_index] = feature.1.into_inner() as f32;
|
|
}
|
|
}
|
|
}
|
|
None => (),
|
|
}
|
|
if self.continuous_features_to_report.contains(feature.0) {
|
|
self.continuous_feature_metrics
|
|
.with_label_values(&[feature.0.to_string().as_str()])
|
|
.observe(feature.1.into_inner())
|
|
} else if self.discrete_features_to_report.contains(feature.0) {
|
|
self.discrete_feature_metrics
|
|
.with_label_values(&[feature.0.to_string().as_str()])
|
|
.observe(feature.1.into_inner())
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process the batch of datarecords
|
|
for r in bpr_start..bpr_end {
|
|
let dr: &DataRecord =
|
|
&bpr.individual_features_list[usize::try_from(r - bpr_start).unwrap()];
|
|
if dr.continuous_features.is_some() {
|
|
for feature in dr.continuous_features.as_ref().unwrap() {
|
|
match self.feature_mapper.get(&feature.0) {
|
|
Some(f_info) => {
|
|
let idx = f_info.index_within_tensor as usize;
|
|
let flat_index: usize = r * cols + idx;
|
|
if flat_index < tensor.len() && idx < cols {
|
|
tensor[flat_index] = feature.1.into_inner() as f32;
|
|
}
|
|
}
|
|
None => (),
|
|
}
|
|
if self.continuous_features_to_report.contains(feature.0) {
|
|
self.continuous_feature_metrics
|
|
.with_label_values(&[feature.0.to_string().as_str()])
|
|
.observe(feature.1.into_inner() as f64)
|
|
} else if self.discrete_features_to_report.contains(feature.0) {
|
|
self.discrete_feature_metrics
|
|
.with_label_values(&[feature.0.to_string().as_str()])
|
|
.observe(feature.1.into_inner() as f64)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
bpr_start = bpr_end;
|
|
}
|
|
|
|
InputTensor::FloatTensor(
|
|
Array2::<f32>::from_shape_vec([rows, cols], tensor)
|
|
.unwrap()
|
|
.into_dyn(),
|
|
)
|
|
}
|
|
|
|
fn get_binary(&self, bprs: &[BatchPredictionRequest], batch_ends: &[usize]) -> InputTensor {
|
|
// These need to be part of model schema
|
|
let rows: usize = batch_ends[batch_ends.len() - 1];
|
|
let cols: usize = 149;
|
|
let full_size: usize = rows * cols;
|
|
let default_val: i64 = 0;
|
|
|
|
let mut v = vec![default_val; full_size];
|
|
|
|
let mut bpr_start = 0;
|
|
for (bpr, &bpr_end) in bprs.iter().zip(batch_ends) {
|
|
// Common features
|
|
if bpr.common_features.is_some()
|
|
&& bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.binary_features
|
|
.is_some()
|
|
{
|
|
let common_features = bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.binary_features
|
|
.as_ref()
|
|
.unwrap();
|
|
|
|
for feature in common_features {
|
|
match self.feature_mapper.get(feature) {
|
|
Some(f_info) => {
|
|
let idx = f_info.index_within_tensor as usize;
|
|
if idx < cols {
|
|
// Set value in each row
|
|
for r in bpr_start..bpr_end {
|
|
let flat_index: usize = r * cols + idx;
|
|
v[flat_index] = 1;
|
|
}
|
|
}
|
|
}
|
|
None => (),
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process the batch of datarecords
|
|
for r in bpr_start..bpr_end {
|
|
let dr: &DataRecord = &bpr.individual_features_list[r - bpr_start];
|
|
if dr.binary_features.is_some() {
|
|
for feature in dr.binary_features.as_ref().unwrap() {
|
|
match self.feature_mapper.get(&feature) {
|
|
Some(f_info) => {
|
|
let idx = f_info.index_within_tensor as usize;
|
|
let flat_index: usize = r * cols + idx;
|
|
v[flat_index] = 1;
|
|
}
|
|
None => (),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
bpr_start = bpr_end;
|
|
}
|
|
InputTensor::Int64Tensor(
|
|
Array2::<i64>::from_shape_vec([rows, cols], v)
|
|
.unwrap()
|
|
.into_dyn(),
|
|
)
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
fn get_discrete(&self, bprs: &[BatchPredictionRequest], batch_ends: &[usize]) -> InputTensor {
|
|
// These need to be part of model schema
|
|
let rows: usize = batch_ends[batch_ends.len() - 1];
|
|
let cols: usize = 320;
|
|
let full_size: usize = rows * cols;
|
|
let default_val: i64 = 0;
|
|
|
|
let mut v = vec![default_val; full_size];
|
|
|
|
let mut bpr_start = 0;
|
|
for (bpr, &bpr_end) in bprs.iter().zip(batch_ends) {
|
|
// Common features
|
|
if bpr.common_features.is_some()
|
|
&& bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.discrete_features
|
|
.is_some()
|
|
{
|
|
let common_features = bpr
|
|
.common_features
|
|
.as_ref()
|
|
.unwrap()
|
|
.discrete_features
|
|
.as_ref()
|
|
.unwrap();
|
|
|
|
for feature in common_features {
|
|
match self.feature_mapper.get(feature.0) {
|
|
Some(f_info) => {
|
|
let idx = f_info.index_within_tensor as usize;
|
|
if idx < cols {
|
|
// Set value in each row
|
|
for r in bpr_start..bpr_end {
|
|
let flat_index: usize = r * cols + idx;
|
|
v[flat_index] = *feature.1;
|
|
}
|
|
}
|
|
}
|
|
None => (),
|
|
}
|
|
if self.discrete_features_to_report.contains(feature.0) {
|
|
self.discrete_feature_metrics
|
|
.with_label_values(&[feature.0.to_string().as_str()])
|
|
.observe(*feature.1 as f64)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process the batch of datarecords
|
|
for r in bpr_start..bpr_end {
|
|
let dr: &DataRecord = &bpr.individual_features_list[usize::try_from(r).unwrap()];
|
|
if dr.discrete_features.is_some() {
|
|
for feature in dr.discrete_features.as_ref().unwrap() {
|
|
match self.feature_mapper.get(&feature.0) {
|
|
Some(f_info) => {
|
|
let idx = f_info.index_within_tensor as usize;
|
|
let flat_index: usize = r * cols + idx;
|
|
if flat_index < v.len() && idx < cols {
|
|
v[flat_index] = *feature.1;
|
|
}
|
|
}
|
|
None => (),
|
|
}
|
|
if self.discrete_features_to_report.contains(feature.0) {
|
|
self.discrete_feature_metrics
|
|
.with_label_values(&[feature.0.to_string().as_str()])
|
|
.observe(*feature.1 as f64)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
bpr_start = bpr_end;
|
|
}
|
|
InputTensor::Int64Tensor(
|
|
Array2::<i64>::from_shape_vec([rows, cols], v)
|
|
.unwrap()
|
|
.into_dyn(),
|
|
)
|
|
}
|
|
|
|
fn get_user_embedding(
|
|
&self,
|
|
bprs: &[BatchPredictionRequest],
|
|
batch_ends: &[usize],
|
|
) -> InputTensor {
|
|
InputTensor::FloatTensor(
|
|
self.get_embedding_tensors(bprs, self.user_embedding_feature_id, batch_ends)
|
|
.into_dyn(),
|
|
)
|
|
}
|
|
|
|
fn get_eng_embedding(
|
|
&self,
|
|
bpr: &[BatchPredictionRequest],
|
|
batch_ends: &[usize],
|
|
) -> InputTensor {
|
|
InputTensor::FloatTensor(
|
|
self.get_embedding_tensors(bpr, self.user_eng_embedding_feature_id, batch_ends)
|
|
.into_dyn(),
|
|
)
|
|
}
|
|
|
|
fn get_author_embedding(
|
|
&self,
|
|
bpr: &[BatchPredictionRequest],
|
|
batch_ends: &[usize],
|
|
) -> InputTensor {
|
|
InputTensor::FloatTensor(
|
|
self.get_embedding_tensors(bpr, self.author_embedding_feature_id, batch_ends)
|
|
.into_dyn(),
|
|
)
|
|
}
|
|
}
|
|
|
|
impl Converter for BatchPredictionRequestToTorchTensorConverter {
|
|
fn convert(&self, batched_bytes: Vec<Vec<u8>>) -> (Vec<InputTensor>, Vec<usize>) {
|
|
let bprs = batched_bytes
|
|
.into_iter()
|
|
.map(|bytes| {
|
|
BatchPredictionRequestToTorchTensorConverter::parse_batch_prediction_request(bytes)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let batch_ends = bprs
|
|
.iter()
|
|
.map(|bpr| bpr.individual_features_list.len())
|
|
.scan(0usize, |acc, e| {
|
|
//running total
|
|
*acc = *acc + e;
|
|
Some(*acc)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
let t1 = self.get_continuous(&bprs, &batch_ends);
|
|
let t2 = self.get_binary(&bprs, &batch_ends);
|
|
//let _t3 = self.get_discrete(&bprs, &batch_ends);
|
|
let t4 = self.get_user_embedding(&bprs, &batch_ends);
|
|
let t5 = self.get_eng_embedding(&bprs, &batch_ends);
|
|
let t6 = self.get_author_embedding(&bprs, &batch_ends);
|
|
|
|
(vec![t1, t2, t4, t5, t6], batch_ends)
|
|
}
|
|
}
|