# 生成测试数据
from sklearn.datasets import make_classification

X, y = make_classification(n_samples=50000, n_features=100, random_state=0)
df = pd.DataFrame(X)
df['label'] = np.random.randint(2,size=(50000))
df1 = spark.createDataFrame(df)
df1 = df1.withColumn('features', Func.array([col(f"{i}") for i in range(0, 100)])).repartition(1000)

# 创建模型并进行预测
%spark2_1.pyspark
import torch.nn as nn

network = nn.Sequential(
    nn.Linear(100, 2560),
    nn.ReLU(),
    nn.Linear(2560, 2560),
    nn.ReLU(),
    nn.Linear(2560, 2)
    #nn.Softmax(dim=1)
)

class network(nn.Module):
    def __init__(self):
        super(network, self).__init__()
        self.l1 = nn.Linear(100, 2560)
        self.l2 = nn.Linear(2560, 2560)
        self.l3 = nn.Linear(2560, 2)
        
    def forward(self, x):
        x = self.l1(x)
        x = self.l2(x)
        x = self.l3(x)
        return x

net = network()
bc_model_state = spark.sparkContext.broadcast(net.state_dict())

def get_model_for_eval():
  # Broadcast the model state_dict
  net.load_state_dict(bc_model_state.value)
  net.eval()
  return net
  
def one_row_predict(x):
    model = get_model_for_eval()
    t = torch.tensor(x, dtype=torch.float32)
    t = model(t).cpu().detach().numpy()
    #prediction = model(t).cpu().detach().item()
    # return prediction
    return list([float(i) for i in t])

one_row_udf = udf(one_row_predict, ArrayType(FloatType()))
df1 = df1.withColumn('pred_one_row', one_row_udf(col('features')))

class Item2vec(nn.Module):
    def __init__(self, cv_dict, csr_cols):
        super(Item2vec, self).__init__()
        pass

    def forward(self, x):
        pass

    def predict(self, x):
        pass

from pyspark import SparkFiles
sc.addFile('/export/models/item2vec.py')
import sys
sys.path.append('/export/models/')

from item2vec import Item2vec

# model 表示训练好的模型
bc_model_state = sc.broadcast(model.state_dict())
net = Item2vec(cv_dict, csr_cols)

def get_model_for_eval_demo():
  # Broadcast the model state_dict
  net.load_state_dict(bc_model_state.value)
  net.eval()
  return net

# udf + withColumn 的方式
def one_row_predict_demo(x)
    x = torch.tensor(x, dtype=torch.float)
    _, prob = bc_model.predict(x)

    return round(float(prob[0]), 4)
    
one_row_predict_demo_udf = udf(one_row_predict_demo, DoubleType())

one_row_predict_demo_udf = udf(one_row_predict_demo, DoubleType())
df = demo.withColumn('demo_prob', one_row_predict_demo_udf('features'))

def one_row_predict_map(rdds):
    bc_model = get_model_for_eval_demo()
    for row in rdds:
        x = torch.tensor(row.x, dtype=torch.float)
        _, prob = bc_model.predict(x)
    
        yield (row['id'], round(float(prob[0]), 4))

df = demo.rdd.mapPartitions(one_row_predict_map).toDF(['id', 'pred_prob'])

# 代码源自 https://github.com/SaeedNajafi/infer-pytorch-pyspark

def basic_row_handler(row):
    return row

def predict_map(index, partition, ml_task,
                batch_size=16,
                row_preprocessor=basic_row_handler,
                row_postprocessor=basic_row_handler):

    # local model loading within each executor
    model = LocalPredictor(ml_task=ml_task, batch_size=batch_size,
                           partition_index=index)

    batch = []
    count = 0
    for row in partition:
        row_dict = row.asDict()
        # apply preprocessor on each row.
        row_dict_prep = row_preprocessor(row_dict)
        batch.append(row_dict_prep)
        count += 1
        if count == batch_size:
            # predict the ml and apply the postprocessor.
            for ret_row in model.predict(batch):  # ml prediction
                ret_row_post = row_postprocessor(ret_row)
                if ret_row_post is not None:
                    yield Row(**ret_row_post)

            batch = []
            count = 0

    # Flush remaining rows in the batches.
    if count != 0:
        for ret_row in model.predict(batch):  # ml prediction
            ret_row_post = row_postprocessor(ret_row)
            if ret_row_post is not None:
                yield Row(**ret_row_post)

        batch = []
        count = 0

# Enable Arrow support.
spark.conf.set("spark.sql.execution.arrow.enabled", "true")
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "64")

sc.addFile('get_model.py')
from get_model import get_model

model_path = '/path/to/model.pt'
data_path = '/path/to/data'

# model 表示训练好的模型
model = torch.load(model_path)
bc_model_state = sc.broadcast(model.state_dict())


def get_model_for_eval():
  # Broadcast the model state_dict  
  model = get_model()
  model.load_state_dict(bc_model_state.value)
  model.eval()
  return model

# model = torch.load(model_path)
# model = sc.broadcast(model)


@pandas_udf(FloatType())
def predict_batch_udf(arr: pd.Series) -> pd.Series:
    model = get_model_for_eval()
    # model.to(device)
    arr = np.vstack(arr.map(lambda x: eval(x)).values)
    arr = torch.tensor(arr).long()
    with torch.no_grad():
        predictions = list(model(arr).cpu().numpy())
            
    return pd.Series(predictions)

# 预测
data = data.withColumn('predictions', predict_batch_udf('features'))