lightgbm_score/applications/functions.py

"""
@author: luojunhui
"""
import jieba.analyse
import pandas as pd

from .model_init import models


class ParamProcess(object):
    """
    处理 params， 继承 models
    """

    def __init__(self):
        self.model_v1 = models.model_v1
        self.model_v2 = models.model_v2
        self.label_encoder = models.label_encoder

    async def title_to_tags(self, features):
        """
        process video title to tags and transform features_json_to_dataFrame
        :param features:
        :return:
        """
        title = features['title']
        if title:
            title = title.strip()
            title_tags = list(jieba.analyse.textrank(title, topK=3))
            if title_tags:
                for i in range(3):
                    try:
                        features['tag{}'.format(i + 1)] = title_tags[i]
                    except:
                        features['tag_{}'.format(i + 1)] = None
            else:
                features['tag1'] = None
                features['tag2'] = None
                features['tag3'] = None
        df = pd.DataFrame([features])
        # print("data_frame", df.columns)
        df = df.drop('title', axis=1)
        return df

    async def predict_score(self, version, features):
        """
        预测
        :param version: 模型版本
        :param features: 视频被 label_encoder 之后的features
        :return: score: 返回的分数
        """
        match version:
            case "v1":
                result = self.model_v1(features)
                result = list(result)
                if result:
                    obj = {
                        "vision": "v1",
                        "score": result[0],
                        "benchmark": 0.06,
                        "is_good_video": 1 if result[0] > 0.06 else 0
                    }
                else:
                    obj = {
                        "vision": "v1",
                        "score": None,
                        "benchmark": 0.06,
                        "is_good_video": 0
                    }
                return obj
            case "v2":
                result = self.model_v2.predict(features)
                result = list(result)
                if result:
                    obj = {
                        "vision": "v2",
                        "score": result[0],
                        "benchmark": 0.3,
                        "is_good_video": 1 if result[0] > 0.3 else 0
                    }
                else:
                    obj = {
                        "vision": "v2",
                        "score": None,
                        "benchmark": 0.3,
                        "is_good_video": 0
                    }
                return obj

    async def process_label(self, params):
        """
        处理类别 features 和 float features
        :param params: 接收到的参数
        :return:
        """
        version = params['version']
        features = params['features']
        features = await self.title_to_tags(features)
        match version:
            case "v1":
                # 全部转化为类别
                str_column = [
                    "channel",
                    "type",
                    "tag1",
                    "tag2",
                    "tag3"
                ]
                for key in str_column:
                    features[key] = self.label_encoder.fit_transform(features[key])
                return version, features
            case "v2":
                float_column = ["out_play_cnt", "out_like_cnt", "out_share_cnt", "lop", "duration"]
                str_column = ["channel", "mode", "out_user_id", "tag1", "tag2", "tag3"]
                for key in float_column:
                    features[key] = pd.to_numeric(features[key], errors="coerce")
                for key in str_column:
                    features[key] = self.label_encoder.fit_transform(features[key])
                return version, features

    async def process(self, params):
        """
        处理
        :param params:
        :return:
        """
        version, features = await self.process_label(params)
        # print(version, features)
        return await self.predict_score(version, features)