Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import torch
 
 # グローバル変数の初期化
@@ -10,30 +10,42 @@ tokenizer = None
 # Hugging Face トークンの取得
 HUGGING_FACE_TOKEN = os.getenv('HUGGINGFACE_TOKEN')
 if not HUGGING_FACE_TOKEN:
-    raise ValueError("環境変数 HUGGINGFACE_TOKEN が設定されていません")
+    raise ValueError("環境変数 HUGGING_FACE_TOKEN が設定されていません")
 
 def load_model():
     global model, tokenizer
     if model is None:
         model_name = "Guchyos/gemma-2b-elyza-task"
         try:
-            # トークンを使用してモデルとトークナイザーを読み込み
+            # まずモデルの設定を読み込む
+            config = AutoConfig.from_pretrained(
+                model_name,
+                token=HUGGING_FACE_TOKEN,
+                trust_remote_code=True
+            )
+            
+            # トークナイザーの読み込み
             tokenizer = AutoTokenizer.from_pretrained(
                 model_name,
-                token=HUGGING_FACE_TOKEN
+                token=HUGGING_FACE_TOKEN,
+                trust_remote_code=True
             )
             
-            # 量子化を無効にし、float32で読み込む設定を追加
+            # モデルの読み込み
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
+                config=config,
                 torch_dtype=torch.float32,
                 device_map="cpu",
                 token=HUGGING_FACE_TOKEN,
-                load_in_8bit=False,      # 8bit量子化を無効化
-                load_in_4bit=False,      # 4bit量子化を無効化
-                quantization_config=None  # 量子化設定を無効化
+                load_in_8bit=False,
+                load_in_4bit=False,
+                trust_remote_code=True
             )
             
+            # モデルを評価モードに設定
+            model.eval()
+            
         except Exception as e:
             raise Exception(f"モデルの読み込みに失敗しました: {str(e)}")
     return model, tokenizer
@@ -48,7 +60,9 @@ def predict(message, history):
             outputs = model.generate(
                 **inputs,
                 max_new_tokens=128,
-                do_sample=False
+                do_sample=False,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id
             )
             
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)