model_prediction_lightgbm.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  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['model'] = model_name
  27. print("model predict result successfully!")
  28. features_reserve = col_reserve + ['model','predict']
  29. return df[set(features_reserve)]
  30. @app.route('/model_prediction_lightgbm', methods=['POST'])
  31. def model_prediction_lightgbm():
  32. # 获取程序开始时间
  33. start_time = time.time()
  34. result = {}
  35. success = 0
  36. print("Program starts execution!")
  37. try:
  38. args = request.values.to_dict()
  39. print('args',args)
  40. logger.info(args)
  41. power_df = get_data_from_mongo(args)
  42. model = model_prediction(power_df,args)
  43. insert_data_into_mongo(model,args)
  44. success = 1
  45. except Exception as e:
  46. my_exception = traceback.format_exc()
  47. my_exception.replace("\n","\t")
  48. result['msg'] = my_exception
  49. end_time = time.time()
  50. result['success'] = success
  51. result['args'] = args
  52. result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
  53. result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
  54. print("Program execution ends!")
  55. return result
  56. if __name__=="__main__":
  57. print("Program starts execution!")
  58. logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
  59. logger = logging.getLogger("model_prediction_lightgbm log")
  60. from waitress import serve
  61. serve(app, host="0.0.0.0", port=10090)
  62. print("server start!")