Update app.py

app.py

@@ -1,65 +1,43 @@
 import os
+import tempfile
 from vllm import LLM
 from vllm.sampling_params import SamplingParams
-from huggingface_hub import hf_hub_download, login
+from huggingface_hub import hf_hub_download
 from datetime import datetime, timedelta
 
-# Hugging Face トークンでログイン
-hf_token = os.getenv("HF_TOKEN")
-if not hf_token:
-    raise ValueError("Hugging Face token is not set in environment variables.")
-login(hf_token)
-
-model_name = "mistralai/Pixtral-Large-Instruct-2411"
+# モデル名
+model_name = "mistralai/Mistral-7B-Instruct-v0.2"
 
+# SYSTEM_PROMPTのロード関数
 def load_system_prompt(repo_id: str, filename: str) -> str:
-    """指定されたリポジトリからSYSTEM_PROMPT.txtをダウンロードし、フォーマット済みプロンプトを返す"""
     file_path = hf_hub_download(repo_id=repo_id, filename=filename)
     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 をロード
-SYSTEM_PROMPT = load_system_prompt(model_name, "SYSTEM_PROMPT.txt")
-
-# 画像URL
-image_url = "https://huggingface.co/datasets/patrickvonplaten/random_img/resolve/main/europe.png"
+# 一時ディレクトリをキャッシュ用に設定
+with tempfile.TemporaryDirectory() as tmpdirname:
+    os.environ["TRANSFORMERS_CACHE"] = tmpdirname
+    os.environ["HF_HOME"] = tmpdirname
+    os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")  # 環境変数からトークンを取得
 
-# メッセージリスト
-messages = [
-    {"role": "system", "content": SYSTEM_PROMPT},
-    {
-        "role": "user",
-        "content": [
-            {
-                "type": "text",
-                "text": "Which of the depicted countries has the best food? Which the second and third and fourth? Name the country, its color on the map and one its city that is visible on the map, but is not the capital. Make absolutely sure to only name a city that can be seen on the map.",
-            },
-            {"type": "image_url", "image_url": {"url": image_url}},
-        ],
-    },
-]
+    SYSTEM_PROMPT = load_system_prompt(model_name, "SYSTEM_PROMPT.txt")
 
-# サンプリング設定
-sampling_params = SamplingParams(max_tokens=512)
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": "Which of the depicted countries has the best food?",
+        },
+    ]
 
-# LLMの初期化 (CPUを利用)
-llm = LLM(
-    model=model_name,
-    config_format="mistral",
-    load_format="mistral",
-    tokenizer_mode="mistral",
-    tensor_parallel_size=1,  # CPUモードでは並列数は1
-    device="cpu",  # 明示的にCPUを指定
-    limit_mm_per_prompt={"image": 4}  # マルチモーダル入力制限
-)
+    sampling_params = SamplingParams(max_tokens=512)
 
-# メッセージを送信し、応答を取得
-outputs = llm.chat(messages, sampling_params=sampling_params)
+    # モデルロードと実行
+    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)
+    print(outputs[0].outputs[0].text)