commit

app.py

@@ -1,96 +1,132 @@
-import gradio as gr
-import torch
-import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
+# import gradio as gr
+# import torch
+# import os
+# from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# 環境変数からトークンを取得
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
-if not HUGGINGFACE_TOKEN:
-    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
+# # 環境変数からトークンを取得
+# HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+# if not HUGGINGFACE_TOKEN:
+#     raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
 
 
-def _load_model():
-    if not torch.cuda.is_available():
-        raise RuntimeError("GPU is not available but required.")
-    print("GPU is available and model will be loaded.")
-    return "GPU ready"
+# def _load_model():
+#     if not torch.cuda.is_available():
+#         raise RuntimeError("GPU is not available but required.")
+#     print("GPU is available and model will be loaded.")
+#     return "GPU ready"
 
-_load_model()
+# _load_model()
 
-# モデルとトークナイザーの初期化
-MODEL_NAME = "google/gemma-7b-it"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HUGGINGFACE_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    token=HUGGINGFACE_TOKEN
-)
-
-def generate_response(prompt):
-    # プロンプトの準備
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+# # モデルとトークナイザーの初期化
+# MODEL_NAME = "google/gemma-7b-it"
+# tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HUGGINGFACE_TOKEN)
+# model = AutoModelForCausalLM.from_pretrained(
+#     MODEL_NAME,
+#     torch_dtype=torch.float16,
+#     device_map="auto",
+#     token=HUGGINGFACE_TOKEN
+# )
+
+# def generate_response(prompt):
+#     # プロンプトの準備
+#     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    # 応答の生成
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
+#     # 応答の生成
+#     with torch.no_grad():
+#         outputs = model.generate(
+#             **inputs,
+#             max_new_tokens=512,
+#             temperature=0.7,
+#             top_p=0.9,
+#             do_sample=True,
+#             pad_token_id=tokenizer.eos_token_id
+#         )
     
-    # 応答のデコード
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+#     # 応答のデコード
+#     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+#     return response
 
-def respond(message, history):
-    # チャット履歴の構築
-    chat_history = ""
-    for msg in history:
-        chat_history += f"{msg['role']}: {msg['content']}\n"
+# def respond(message, history):
+#     # チャット履歴の構築
+#     chat_history = ""
+#     for msg in history:
+#         chat_history += f"{msg['role']}: {msg['content']}\n"
     
-    # 現在のメッセージを追加
-    prompt = f"{chat_history}Human: {message}\nAssistant:"
+#     # 現在のメッセージを追加
+#     prompt = f"{chat_history}Human: {message}\nAssistant:"
     
-    try:
-        response = generate_response(prompt)
-        # 応答から余分な部分を削除
-        response = response.split("Assistant:")[-1].strip()
-        return response
-    except Exception as e:
-        return f"エラーが発生しました: {str(e)}"
-
-# Gradioインターフェースの設定
-iface = gr.ChatInterface(
-    fn=respond,
-    textbox=gr.Textbox(
-        placeholder="メッセージを入力してください...",
-        container=False,
-        scale=7,
-        lines=2
-    ),
-    chatbot=gr.Chatbot(
-        height=600,
-        show_copy_button=True,
-        show_share_button=True,
-        avatar_images=(None, None)
-    ),
-    title="Gemma Chat Assistant",
-    description="Google Gemmaモデルを使用したチャットアシスタントです。",
-    theme=gr.themes.Soft(),
-    examples=[
-        "こんにちは",
-        "自己紹介をしてください",
-        "Pythonについて教えてください"
-    ]
-)
+#     try:
+#         response = generate_response(prompt)
+#         # 応答から余分な部分を削除
+#         response = response.split("Assistant:")[-1].strip()
+#         return response
+#     except Exception as e:
+#         return f"エラーが発生しました: {str(e)}"
+
+# # Gradioインターフェースの設定
+# iface = gr.ChatInterface(
+#     fn=respond,
+#     textbox=gr.Textbox(
+#         placeholder="メッセージを入力してください...",
+#         container=False,
+#         scale=7,
+#         lines=2
+#     ),
+#     chatbot=gr.Chatbot(
+#         height=600,
+#         show_copy_button=True,
+#         show_share_button=True,
+#         avatar_images=(None, None)
+#     ),
+#     title="Gemma Chat Assistant",
+#     description="Google Gemmaモデルを使用したチャットアシスタントです。",
+#     theme=gr.themes.Soft(),
+#     examples=[
+#         "こんにちは",
+#         "自己紹介をしてください",
+#         "Pythonについて教えてください"
+#     ]
+# )
+
+# if __name__ == "__main__":
+#     iface.launch(
+#         share=True,
+#         server_name="0.0.0.0",
+#         server_port=7860
+#     ) 
+
+
+import os, torch, gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+MODEL_NAME = "google/gemma-7b-it"
+
+model, tokenizer = None, None        # ← グローバルで空のまま
+
+def load_model():
+    """初回リクエスト時にのみ GPU を要求してモデルをロード"""
+    global model, tokenizer
+    if model is not None:
+        return
+    if not torch.cuda.is_available():
+        # ZeroGPU ならここで一度 False → 数秒待って再度 True になることもある
+        raise RuntimeError("GPU still not attached (ZeroGPU)。数秒後に再試行してください。")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HUGGINGFACE_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        token=HUGGINGFACE_TOKEN
+    )
+
+def respond(message, history):
+    load_model()                               # ← ここで初めて GPU を確保・モデルロード
+    inputs = tokenizer(message, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        out = model.generate(**inputs, max_new_tokens=512, temperature=0.7, top_p=0.9)
+    return tokenizer.decode(out[0], skip_special_tokens=True)
 
+iface = gr.ChatInterface(fn=respond, title="Gemma-ZeroGPU Demo")
 if __name__ == "__main__":
-    iface.launch(
-        share=True,
-        server_name="0.0.0.0",
-        server_port=7860
-    ) 
+    iface.launch()