config.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. # put your config here
  2. import argparse, os, sys
  3. from funasr import AutoModel
  4. from dotenv import load_dotenv
  5. def load_env_from_single_arg():
  6. """
  7. 通过命令行参数 --env <路径> 加载环境变量文件。
  8. 若不传 --env,则不加载。
  9. """
  10. if "--env" in sys.argv:
  11. try:
  12. idx = sys.argv.index("--env")
  13. env_path = sys.argv[idx + 1]
  14. load_dotenv(dotenv_path=env_path)
  15. print(f"✔ 加载环境变量文件: {env_path}")
  16. return True, env_path
  17. except IndexError:
  18. print("❌ 错误:--env 后缺少路径")
  19. else:
  20. print("⚠ 未传 --env,未加载任何环境变量")
  21. return False, None
  22. load_env_from_single_arg()
  23. port = int(os.getenv('PORT',8067))
  24. file_url = os.getenv('FILE_URL',f"http://10.41.1.220:{port}")
  25. api_key = os.getenv('API_KEY')
  26. base_url = os.getenv('BASE_URL','https://dashscope.aliyuncs.com/compatible-mode/v1')
  27. llm_config={
  28. "config_list": [
  29. {
  30. "model": os.getenv('MODEL', 'qwen-max'), # Same as in vLLM command
  31. "api_key": api_key, # Not needed
  32. "base_url": base_url # Your vLLM URL, with '/v1' added
  33. }
  34. ],
  35. "cache_seed": None, # Turns off caching, useful for testing different models
  36. "temperature": float(os.getenv('MODEL_TEM_TURE',0.5)),
  37. }
  38. llm_config_ds={
  39. "config_list": [
  40. {
  41. "model": os.getenv('THINK_MODEL', 'deepseek-r1'), # Same as in vLLM command
  42. "api_key": api_key, # Not needed
  43. "base_url": base_url # Your vLLM URL, with '/v1' added
  44. }
  45. ],
  46. "cache_seed": None, # Turns off caching, useful for testing different models
  47. "temperature": float(os.getenv('THINK_MODEL_TEM_TURE',0.7)),
  48. }
  49. llm_config_plus={
  50. "config_list": [
  51. {
  52. "model": os.getenv('PLUS_MODEL', 'qwen-plus'), # Same as in vLLM command
  53. "api_key": api_key, # Not needed
  54. "base_url": base_url # Your vLLM URL, with '/v1' added
  55. }
  56. ],
  57. "cache_seed": None, # Turns off caching, useful for testing different models
  58. "temperature": float(os.getenv('PLUS_MODEL_TEM_TURE',0.7)),
  59. }
  60. llm_config_turbo={
  61. "config_list": [
  62. {
  63. "model": os.getenv('TURBO_MODEL', 'qwen-plus'), # Same as in vLLM command
  64. "api_key": api_key, # Not needed
  65. "base_url": base_url # Your vLLM URL, with '/v1' added
  66. }
  67. ],
  68. "cache_seed": None, # Turns off caching, useful for testing different models
  69. "temperature": float(os.getenv('TURBO_MODEL_TEM_TURE',0.7)),
  70. }
  71. llm_dict = {'deepseek':llm_config_ds, 'turbo':llm_config_turbo, 'plus':llm_config_plus,'max':llm_config,}
  72. STATIC_DIR = os.getenv('STATIC_DIR','/workspace')
  73. BASE_UPLOAD_DIRECTORY = os.getenv('BASE_UPLOAD_DIRECTORY','/app/workspace')
  74. db_list_2 = [{
  75. "host": "10.41.1.220",
  76. "database": "sale_database",
  77. "user": "sa",
  78. "password": "wangdalin@666",
  79. "port": "1433"
  80. },
  81. {
  82. "host": "10.41.1.220",
  83. "database": "cooperation",
  84. "user": "sa",
  85. "password": "wangdalin@666",
  86. "port": "1433",
  87. }]
  88. db_cn_map = {'sale_database':{
  89. 'name':'销售数据库','table_cn_map':{
  90. 'sale':'销售数据'}},
  91. 'cooperation':{'name':'商业往来数据库','table_cn_map':{
  92. 'SupplierPurchases':'供应商数据'}}}
  93. def reverse_cn_map(cn_map):
  94. # 创建中文找英文的结构
  95. en_to_cn_map = {}
  96. for db_key, db_value in cn_map.items():
  97. db_name_cn = db_value.get('name', db_key)
  98. en_to_cn_map[db_name_cn] = {
  99. 'name': db_key,
  100. 'table_cn_map': {}
  101. }
  102. for table_key, table_value in db_value['table_cn_map'].items():
  103. table_name_cn = table_value
  104. en_to_cn_map[db_name_cn]['table_cn_map'][table_name_cn] = table_key
  105. return en_to_cn_map
  106. db_en_map = reverse_cn_map(db_cn_map)
  107. stt_model_dir = os.getenv('STT_MODEL', "/model/paraformer-zh")
  108. stt_device = os.getenv('STT_DEVICE', "cpu")
  109. stt_model = AutoModel(
  110. model=stt_model_dir,
  111. vad_model=os.getenv('VAD_MODEL', "/model/fsmn-vad"),
  112. vad_kwargs={"max_single_segment_time": 30000},
  113. punc_model=os.getenv('PUNC_MODEL', "/model/ct-punc"),
  114. device="cpu",
  115. )
  116. bge_model_path = os.getenv('BGE_MODEL', '/model/bge-reranker-v2-m3',)
  117. if __name__ == '__main__':
  118. db_en_map = reverse_cn_map(db_cn_map)
  119. print(db_en_map)