model_prediction_lightgbm.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. import pandas as pd
  2. from pymongo import MongoClient
  3. import pickle
  4. from flask import Flask,request
  5. import time
  6. import logging
  7. import traceback
  8. from common.database_dml import get_data_from_mongo,insert_data_into_mongo
  9. app = Flask('model_prediction_lightgbm——service')
  10. def str_to_list(arg):
  11. if arg == '':
  12. return []
  13. else:
  14. return arg.split(',')
  15. def model_prediction(df,args):
  16. mongodb_connection,mongodb_database,mongodb_model_table,model_name,col_reserve = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",args['mongodb_database'],args['mongodb_model_table'],args['model_name'],str_to_list(args['col_reserve'])
  17. client = MongoClient(mongodb_connection)
  18. db = client[mongodb_database]
  19. collection = db[mongodb_model_table]
  20. model_data = collection.find_one({"model_name": model_name})
  21. if model_data is not None:
  22. model_binary = model_data['model'] # 确保这个字段是存储模型的二进制数据
  23. # 反序列化模型
  24. model = pickle.loads(model_binary)
  25. df['predict'] = model.predict(df[model.feature_name()])
  26. df.loc[df['predict']<0,'predict']=0
  27. df['model'] = model_name
  28. print("model predict result successfully!")
  29. features_reserve = col_reserve + ['model','predict']
  30. return df[set(features_reserve)]
  31. @app.route('/model_prediction_lightgbm', methods=['POST'])
  32. def model_prediction_lightgbm():
  33. # 获取程序开始时间
  34. start_time = time.time()
  35. result = {}
  36. success = 0
  37. print("Program starts execution!")
  38. try:
  39. args = request.values.to_dict()
  40. print('args',args)
  41. logger.info(args)
  42. power_df = get_data_from_mongo(args)
  43. model = model_prediction(power_df,args)
  44. insert_data_into_mongo(model,args)
  45. success = 1
  46. except Exception as e:
  47. my_exception = traceback.format_exc()
  48. my_exception.replace("\n","\t")
  49. result['msg'] = my_exception
  50. end_time = time.time()
  51. result['success'] = success
  52. result['args'] = args
  53. result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
  54. result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
  55. print("Program execution ends!")
  56. return result
  57. if __name__=="__main__":
  58. print("Program starts execution!")
  59. logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
  60. logger = logging.getLogger("model_prediction_lightgbm log")
  61. from waitress import serve
  62. serve(app, host="0.0.0.0", port=10090)
  63. print("server start!")