app.py

import os
from vllm import LLM
from vllm.sampling_params import SamplingParams
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},
    {
        "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}},
        ],
    },
]

# サンプリング設定
sampling_params = SamplingParams(max_tokens=512)

# 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)