feat: add env config

app.py

@@ -5,8 +5,13 @@ import gradio as gr
 from huggingface_hub import login
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+import env_config
+
+api_token = env_config.HF_API_TOKEN
+max_new_tokens = env_config.MAX_NEW_TOKENS
+model_id = env_config.MODEL_ID
+
 # 登录 Hugging Face API
-api_token = os.environ.get("HF_API_TOKEN")
 login(api_token)
 
 # 模型加载函数
@@ -19,7 +24,7 @@ def get_llm(model_id):
 @spaces.GPU(duration=120)
 def retriever_qa(file, query):
     # 加载模型和分词器
-    model_id = 'mistralai/Mistral-7B-Instruct-v0.2'
+    # model_id = 'mistralai/Mistral-7B-Instruct-v0.2'
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
 
     # 确保 CUDA 初始化不在主线程
@@ -47,7 +52,7 @@ def retriever_qa(file, query):
         print('Start Inference')
         
         # 推理
-        generated_ids = llm.generate(model_inputs, max_new_tokens=50, do_sample=True)
+        generated_ids = llm.generate(model_inputs, max_new_tokens=max_new_tokens, do_sample=True)
         # generated_ids = llm.generate(input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask'], max_new_tokens=50, do_sample=True)
         print(f'Generated ids: {generated_ids}')
         # 解码输出

env_config.py

@@ -0,0 +1,5 @@
+import os
+
+HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
+MAX_NEW_TOKENS = os.environ.get("MAX_NEW_TOKENS", 1024)
+MODEL_ID = os.environ.get("MODEL_ID", "mistralai/Mistral-7B-Instruct-v0.2")