commit

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 from typing import List, Tuple
+import spaces
 
 # Hugging Face token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -12,50 +13,54 @@ class ChatBot:
         self.model = None
         self.tokenizer = None
         self.current_model = None
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
     def load_model(self, model_name: str):
         """モデルとトークナイザーをロード"""
-        if self.current_model == model_name:
+        if self.current_model == model_name and self.model is not None:
             return
             
         # メモリクリア
         if self.model is not None:
             del self.model
-            del self.tokenizer
             torch.cuda.empty_cache()
         
-        # モデルロード
+        # トークナイザーロード
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_name,
             use_auth_token=HF_TOKEN,
             trust_remote_code=True
         )
         
+        # パッドトークンの設定
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        # モデルロード（ZeroGPU対応）
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name,
             use_auth_token=HF_TOKEN,
-            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-            device_map="auto" if self.device == "cuda" else None,
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
             trust_remote_code=True
         )
-        
-        if self.device == "cuda":
-            self.model = self.model.to(self.device)
             
         self.current_model = model_name
         
+    @spaces.GPU(duration=60)
     def generate_response(self, message: str, history: List[Tuple[str, str]], model_name: str, 
                          temperature: float = 0.7, max_tokens: int = 512) -> str:
         """メッセージに対する応答を生成"""
         # モデルロード
         self.load_model(model_name)
         
+        # GPUに移動
+        self.model.to('cuda')
+        
         # プロンプト構築
         prompt = self._build_prompt(message, history)
         
         # トークナイズ
-        inputs = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+        inputs = self.tokenizer.encode(prompt, return_tensors="pt").to('cuda')
         
         # 生成
         with torch.no_grad():
@@ -71,6 +76,11 @@ class ChatBot:
         
         # デコード
         response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+        
+        # CPUに戻す（メモリ節約）
+        self.model.to('cpu')
+        torch.cuda.empty_cache()
+        
         return response.strip()
         
     def _build_prompt(self, message: str, history: List[Tuple[str, str]]) -> str:
@@ -143,7 +153,7 @@ with gr.Blocks(title="ChatGPT Clone", theme=gr.themes.Soft()) as app:
             model_select = gr.Dropdown(
                 choices=[
                     "elyza/Llama-3-ELYZA-JP-8B",
-                    "Fugaku-LLM/Fugaku-LLM-13B"
+                    "cyberagent/open-calm-7b"
                 ],
                 value="elyza/Llama-3-ELYZA-JP-8B",
                 label="モデル選択",
@@ -176,7 +186,8 @@ with gr.Blocks(title="ChatGPT Clone", theme=gr.themes.Soft()) as app:
             
             ### 注意事項
             - 初回のモデル読み込みには時間がかかります
-            - GPU使用時は高速に動作します
+            - ZeroGPU使用により高速推論が可能
+            - 1回の生成は60秒以内に完了します
             """)
     
     # イベントハンドラ