Update app.py

app.py

@@ -2,30 +2,30 @@ import os
 import tempfile
 from vllm import LLM
 from vllm.sampling_params import SamplingParams
-from huggingface_hub import hf_hub_download
 from datetime import datetime, timedelta
 
 # モデル名
 model_name = "mistralai/Mistral-7B-Instruct-v0.2"
 
-# SYSTEM_PROMPTのロード関数
-def load_system_prompt(repo_id: str, filename: str) -> str:
-    file_path = hf_hub_download(repo_id=repo_id, filename=filename)
+# SYSTEM_PROMPTのロード関数 (ローカルファイルを読み込む)
+def load_system_prompt(file_path: str) -> str:
     with open(file_path, 'r') as file:
         system_prompt = file.read()
     today = datetime.today().strftime('%Y-%m-%d')
     yesterday = (datetime.today() - timedelta(days=1)).strftime('%Y-%m-%d')
-    model_name = repo_id.split("/")[-1]
     return system_prompt.format(name=model_name, today=today, yesterday=yesterday)
 
+# SYSTEM_PROMPT.txtは現在のディレクトリ内にあることを想定
+system_prompt_path = "./SYSTEM_PROMPT.txt"  # 現在のディレクトリ内のファイルパス
+SYSTEM_PROMPT = load_system_prompt(system_prompt_path)
+
 # 一時ディレクトリをキャッシュ用に設定
 with tempfile.TemporaryDirectory() as tmpdirname:
     os.environ["TRANSFORMERS_CACHE"] = tmpdirname
     os.environ["HF_HOME"] = tmpdirname
     os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")  # 環境変数からトークンを取得
 
-    SYSTEM_PROMPT = load_system_prompt(model_name, "SYSTEM_PROMPT.txt")
-
+    # メッセージリストの作成
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {
@@ -40,4 +40,5 @@ with tempfile.TemporaryDirectory() as tmpdirname:
     llm = LLM(model=model_name, trust_remote_code=True, tensor_parallel_size=1, device="cpu")
     outputs = llm.chat(messages, sampling_params=sampling_params)
 
+    # 出力を表示
     print(outputs[0].outputs[0].text)