main.py

# -*- coding: utf-8 -*
# from __future__ import print_function
import sys
import tensorflow as tf
# import tensorflow_datasets as tfds
import numpy as np
import json

tf.enable_eager_execution()

def test():
    # mirrored_strategy = tf.distribute.MirroredStrategy()
    # # 在config中加入镜像策略
    # config = tf.estimator.RunConfig(train_distribute=mirrored_strategy, eval_distribute=mirrored_strategy)
    # 把config加到模型里
    regressor = tf.estimator.LinearRegressor(
        feature_columns=[tf.feature_column.numeric_column('feats')],
        optimizer='SGD'
        # ,config=config
    )
    def input_fn():
        dataset = tf.data.Dataset.from_tensors(({"feats":[1.]}, [1.]))
        return dataset.repeat(1000).batch(10)

    # 正常训练，正常评估
    regressor.train(input_fn=input_fn
                    , steps=20
                    )
    regressor.evaluate(input_fn=input_fn
                       # , steps=10
                       )

def parse_from_json(config_path):
    """ parse feature columns from feature config path

    Args:
      config_path: string, a feature config path
    """
    total = 0
    correct = 0
    with open(config_path, "r") as f:
        config = json.load(f)

    feature_names = set()
    features = config["features"]
    for feature in features:
        feature_name = feature['feature_name']
        if '#' in feature_name:
            feature_name = feature_name.split('#')[0]
        feature_names.add(feature_name)
    return feature_names

#convert model's format from *.pb to *.pbtxt
def parse_model_2_txt(saved_model_dir ,output_file):
    from tensorflow.python.saved_model import loader_impl
    from google.protobuf import text_format
    saved_model = loader_impl._parse_saved_model(saved_model_dir)
    with open(output_file, 'w') as f:
        f.write(text_format.MessageToString(saved_model))

# parse_model_2_txt('/Users/machi/git/internal/starship_galaxy/model_zoo/scheduler/2022q2combo/old', '/Users/machi/git/internal/starship_galaxy/model_zoo/scheduler/2022q2combo/old/saved_model.pbtxt')

import os
def build_serving_input_new():
    import pickle
    with tf.gfile.Open('feature_desc.pkl', mode='rb') as f:
        feature_dec = pickle.load(f)
    sep_placeholder = {}
    for name, desc in feature_dec.items():
        if 'sg_poi_click_time_gap_seq_2d' in name:
            print(desc)

    # return sep_placeholder

def read_schema(file):
    d = {}
    with open(file) as f:
        for line in f:
            line = line.strip()
            fds = line.split(' ')
            d[fds[0]] = fds[1]
    return d



def sparse_tensor():
    indices_tf = tf.constant([[0, 0], [0, 1], [1, 1], [2, 2]], dtype=tf.int64)
    values_tf = tf.constant([1, 2, 3, 4], dtype=tf.int32)
    dense_shape_tf = tf.constant([3, 3], dtype=tf.int64)

    sparse_tf = tf.SparseTensor(indices=indices_tf,
                                values=values_tf,
                                dense_shape=dense_shape_tf)
    dense_tf = tf.sparse_tensor_to_dense(sparse_tf)

    # print(dense_tf)


    user_tf = tf.constant([1, 2, 3], dtype=tf.int32, shape=[3, 1])

    # 一行为一个session，每一行包含不同个数的样本。以下示例中，共有3个session，第1个session包含3个样本，第2个session包含2个样本，第3个session行包含1个样本
    # b为non_common特征
    b = tf.constant([[1, 2, 1], [0, 3, 2], [0, 0, 4]])

    # a为common特征，3个session有3个值
    a = tf.constant([1, 2, 3], shape=[3, 1])

    # 将a扩展为和b相同维度
    a = tf.tile(a, tf.constant([1, 3]))
    print(a)

    # 获取b中非0元素的下标
    indices = tf.where(tf.not_equal(b, 0))
    print(indices)

    # 将非0元素的下标处的a和b的值拼接起来，即样本展开后的结果
    c = tf.concat(values=[tf.expand_dims(tf.gather_nd(a, indices), axis=1), tf.expand_dims(tf.gather_nd(b, indices), axis=1)], axis=1)
    print(c)


def kkv_attention(query, key, value, mask=None):
    # Transpose key and value matrices
    key_transpose = tf.transpose(key, perm=[0, 2, 1])
    value_transpose = tf.transpose(value, perm=[0, 2, 1])

    # Compute dot product between query and key
    logits = tf.matmul(query, key_transpose)

    # Apply mask (if provided) to logits
    if mask is not None:
        logits += mask

    # Apply softmax activation to obtain attention scores
    attention_scores = tf.nn.softmax(logits, axis=-1)

    # Apply attention scores to value to obtain context vector
    context_vector = tf.matmul(attention_scores, value_transpose)

    # Transpose back the output
    context_vector = tf.transpose(context_vector, perm=[0, 2, 1])

    return context_vector, attention_scores

# write kkv attention function
def write_kkv_attention(query, key, value, mask=None):
    # Transpose key and value matrices
    # key_transpose = tf.transpose(key, perm=[0, 2, 1])
    # value_transpose = tf.transpose(value, perm=[0, 2, 1])

    # Compute dot product between query and key
    logits = tf.matmul(query, key)

    # Apply mask (if provided) to logits
    if mask is not None:
        logits += mask

    # Apply softmax activation to obtain attention scores
    attention_scores =   tf.nn.softmax(logits, axis=-1)

    # Apply attention scores to value to obtain context vector
    context_vector =  tf.matmul(attention_scores, value)

    # Transpose back the output
    # context_vector =  tf.transpose(context_vector, perm=[0, 2, 1])

    return context_vector, attention_scores

# test write_kkv_attention
def test_write_kkv_attention():
    # define query and key matrices
    query =  tf.constant([[-0.1250,  0.0000, -0.5000,  0.5000,  0.0000]])

    key =  tf.constant([[ -0.1250,  0.0000, -0.5000,  0.5000,  0.0000],
                        [-0.5000,  0.0000,  0.5000,  0.5000,  0.0000],
                        [-0.2500, -0.5000,  0.0000,  0.5000,  0.2500],
                        [ 0.0000,  0.0000,  0.0000,  0.5000,  0.5000],
                        [ 0.5000,  0.5000,  0.0000, -0.5000,  0.5000]])

    value =   tf.constant([[-0.5000,  0.0000,  0.5000,  0.5000,  0.0000],
                        [-0.5000,  0.0000,  0.5000,  0.5000,  0.0000],
                        [-0.5000,  0.0000,  0.5000,  0.5000,  0.0000],
                        [ 0.0000,  0.0000,  0.5000,  0.5000,  0.5000],
                        [ 0.5000,  0.5000,  0.0000, -0.5000,  0.5000]])


    mask = None

    # call write_kkv_attention and obtain context vector and attention scores
    context_vector, attention_scores = write_kkv_attention(query, key, value,mask)

    # print results
    print context_vector
    print attention_scores


print '123', 1