Update app.py

app.py

@@ -1,36 +1,34 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import os
-from huggingface_hub import login
 
-token = os.environ.get("HUGGINGFACE_TOKEN")
-if token:
-    login(token=token)
+model = None
+tokenizer = None
 
-def predict(message, history):
-    try:
+def load_model():
+    global model, tokenizer
+    if model is None:
         model_name = "Guchyos/gemma-2b-elyza-task"
-        tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
-        
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            device_map="cpu",
-            load_in_8bit=False,  # 8bit量子化を無効化
-            load_in_4bit=False,  # 4bit量子化を無効化
-            use_auth_token=True
+            torch_dtype=torch.float32,  # float32を使用
+            device_map="cpu"
         )
-        
+    return model, tokenizer
+
+def predict(message, history):
+    try:
+        model, tokenizer = load_model()
         prompt = f"質問: {message}\n\n回答:"
-        inputs = tokenizer(prompt, return_tensors="pt")  # .to(device)を削除
+        inputs = tokenizer(prompt, return_tensors="pt")
         
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,  # トークン数を減らして軽量化
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=128,
+                do_sample=False
+            )
         
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response.replace(prompt, "").strip()
@@ -40,25 +38,9 @@ def predict(message, history):
 
 demo = gr.ChatInterface(
     fn=predict,
-    title="💬 Gemma 2 Quantized for ELYZA-tasks",
-    description="""
-    # ELYZA-tasks-100-TV用に最適化された日本語LLMです
-    
-    ## 使い方
-    - 質問を入力してEnterキーを押してください
-    - 生成には数秒かかります
-    
-    ## 特徴
-    - CPU対応
-    - 日本語に特化
-    - ELYZA-tasks形式に対応
-    """,
-    examples=[
-        "日本の四季について、それぞれの特徴を説明してください。",
-        "人工知能の発展における倫理的な課題について説明してください。",
-        "東京の主要な観光スポットを3つ挙げて、それぞれ説明してください。"
-    ]
+    title="💬 Gemma 2 for ELYZA-tasks",
+    description="ELYZA-tasks-100-TV用に最適化された日本語LLMです"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)