commit

app.py

@@ -3,11 +3,17 @@ 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")
 
+# Check if running on ZeroGPU
+try:
+    import spaces
+    IS_ZEROGPU = True
+except ImportError:
+    IS_ZEROGPU = False
+
 class ChatBot:
     def __init__(self):
         self.model = None
@@ -46,10 +52,9 @@ class ChatBot:
             
         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:
-        """メッセージに対する応答を生成"""
+    def _generate_response_gpu(self, message: str, history: List[Tuple[str, str]], model_name: str, 
+                             temperature: float = 0.7, max_tokens: int = 512) -> str:
+        """GPU上で応答を生成する実際の処理"""
         # モデルロード
         self.load_model(model_name)
         
@@ -82,6 +87,37 @@ class ChatBot:
         torch.cuda.empty_cache()
         
         return response.strip()
+    
+    def generate_response(self, message: str, history: List[Tuple[str, str]], model_name: str, 
+                         temperature: float = 0.7, max_tokens: int = 512) -> str:
+        """メッセージに対する応答を生成"""
+        if IS_ZEROGPU:
+            # ZeroGPU環境の場合
+            return self._generate_response_gpu(message, history, model_name, temperature, max_tokens)
+        else:
+            # 通常環境の場合
+            self.load_model(model_name)
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+            
+            if device == 'cuda':
+                self.model.to(device)
+            
+            prompt = self._build_prompt(message, history)
+            inputs = self.tokenizer.encode(prompt, return_tensors="pt").to(device)
+            
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs,
+                    max_new_tokens=max_tokens,
+                    temperature=temperature,
+                    do_sample=True,
+                    top_p=0.95,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
+                )
+            
+            response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+            return response.strip()
         
     def _build_prompt(self, message: str, history: List[Tuple[str, str]]) -> str:
         """会話履歴からプロンプトを構築"""
@@ -99,6 +135,10 @@ class ChatBot:
 # ChatBotインスタンス
 chatbot = ChatBot()
 
+# ZeroGPU環境の場合、GPUデコレータを適用
+if IS_ZEROGPU:
+    chatbot._generate_response_gpu = spaces.GPU(duration=60)(chatbot._generate_response_gpu)
+
 def respond(message: str, history: List[Tuple[str, str]], model_name: str, 
             temperature: float, max_tokens: int) -> Tuple[List[Tuple[str, str]], str]:
     """Gradioのコールバック関数"""
@@ -152,10 +192,10 @@ with gr.Blocks(title="ChatGPT Clone", theme=gr.themes.Soft()) as app:
         with gr.Column(scale=1):
             model_select = gr.Dropdown(
                 choices=[
-                    "elyza/Llama-3-ELYZA-JP-8B",
-                    "cyberagent/open-calm-7b"
+                    "rinna/japanese-gpt2-medium",
+                    "cyberagent/open-calm-small"
                 ],
-                value="elyza/Llama-3-ELYZA-JP-8B",
+                value="rinna/japanese-gpt2-medium",
                 label="モデル選択",
                 interactive=True
             )
@@ -209,4 +249,7 @@ with gr.Blocks(title="ChatGPT Clone", theme=gr.themes.Soft()) as app:
     )
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch(
+        share=False,
+        show_error=True
+    )