Update app.py

app.py

@@ -1,23 +1,35 @@
+import os
 from vllm import LLM
 from vllm.sampling_params import SamplingParams
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, login
 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"
 
 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"
 
+# メッセージリスト
 messages = [
     {"role": "system", "content": SYSTEM_PROMPT},
     {
@@ -32,11 +44,21 @@ messages = [
     },
 ]
 
+# サンプリング設定
 sampling_params = SamplingParams(max_tokens=512)
 
-# note that running this model on GPU requires over 300 GB of GPU RAM
-llm = LLM(model=model_name, config_format="mistral", load_format="mistral", tokenizer_mode="mistral", tensor_parallel_size=8, limit_mm_per_prompt={"image": 4})
+# LLMの初期化 (GPUを利用)
+llm = LLM(
+    model=model_name,
+    config_format="mistral",
+    load_format="mistral",
+    tokenizer_mode="mistral",
+    tensor_parallel_size=8,  # GPUを8個使用
+    limit_mm_per_prompt={"image": 4}  # マルチモーダル入力制限
+)
 
+# メッセージを送信し、応答を取得
 outputs = llm.chat(messages, sampling_params=sampling_params)
 
+# 結果を表示
 print(outputs[0].outputs[0].text)