Create app.py

app.py

@@ -0,0 +1,180 @@
+import torch
+import time
+import uuid
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import List, Optional, Union, Dict, Any
+import uvicorn
+
+app = FastAPI()
+
+# CORS設定
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# INT4量子化設定
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+
+# モデルとトークナイザーの初期化
+model_name = "Qwen/Qwen3-235B-A22B"
+model_id = "qwen3-235b-a22b-int4"  # OpenAI API互換のモデルID
+
+print(f"Loading model: {model_name}")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+print("Model loaded successfully")
+
+# OpenAI API互換のデータモデル
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 512
+    stream: Optional[bool] = False
+
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: ChatMessage
+    finish_reason: str
+
+class ChatCompletionUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: ChatCompletionUsage
+
+class ModelInfo(BaseModel):
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str = "huggingface"
+
+class ModelsResponse(BaseModel):
+    object: str = "list"
+    data: List[ModelInfo]
+
+# チャット形式のプロンプト構築
+def build_chat_prompt(messages: List[ChatMessage]) -> str:
+    """Qwen形式のチャットプロンプトを構築"""
+    prompt = ""
+    for msg in messages:
+        if msg.role == "system":
+            prompt += f"System: {msg.content}\n\n"
+        elif msg.role == "user":
+            prompt += f"Human: {msg.content}\n\n"
+        elif msg.role == "assistant":
+            prompt += f"Assistant: {msg.content}\n\n"
+    
+    # 最後にAssistantのプロンプトを追加
+    if messages[-1].role != "assistant":
+        prompt += "Assistant: "
+    
+    return prompt
+
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def chat_completions(request: ChatCompletionRequest):
+    """OpenAI互換のチャット補完エンドポイント"""
+    try:
+        # プロンプト構築
+        prompt = build_chat_prompt(request.messages)
+        
+        # トークナイズ
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        input_length = inputs["input_ids"].shape[1]
+        
+        # 生成
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # デコード
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # プロンプトを除去
+        response_text = generated_text[len(prompt):].strip()
+        
+        # トークン数計算
+        output_length = outputs[0].shape[0] - input_length
+        
+        # レスポンス構築
+        response = ChatCompletionResponse(
+            id=f"chatcmpl-{uuid.uuid4().hex[:8]}",
+            created=int(time.time()),
+            model=request.model,
+            choices=[
+                ChatCompletionChoice(
+                    index=0,
+                    message=ChatMessage(role="assistant", content=response_text),
+                    finish_reason="stop"
+                )
+            ],
+            usage=ChatCompletionUsage(
+                prompt_tokens=input_length,
+                completion_tokens=output_length,
+                total_tokens=input_length + output_length
+            )
+        )
+        
+        return response
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/v1/models", response_model=ModelsResponse)
+async def list_models():
+    """利用可能なモデルのリストを返す"""
+    return ModelsResponse(
+        data=[
+            ModelInfo(
+                id=model_id,
+                created=int(time.time())
+            )
+        ]
+    )
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "healthy",
+        "model": model_name,
+        "model_id": model_id,
+        "quantization": "INT4"
+    }
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)