Update app.py

app.py

@@ -8,36 +8,39 @@ tokenizer = None
 def load_model():
     global model, tokenizer
     if model is None:
-        model_name = "line-corporation/japanese-large-lm-1.7b"  # LINEの軽量モデル
+        model_name = "Guchyos/gemma-2b-elyza-task"
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            device_map="cpu",
-            low_cpu_mem_usage=True
+            torch_dtype=torch.float32,  # float32を使用
+            device_map="cpu"
         )
     return model, tokenizer
 
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-import sentencepiece as spm
-
 def predict(message, history):
     try:
-        model_name = "rinna/japanese-gpt-neox-small"
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            device_map="cpu",
-            trust_remote_code=True
-        )
+        model, tokenizer = load_model()
+        prompt = f"質問: {message}\n\n回答:"
+        inputs = tokenizer(prompt, return_tensors="pt")
         
-        inputs = tokenizer(message, return_tensors="pt")
-        outputs = model.generate(**inputs, max_length=64)
-        return tokenizer.decode(outputs[0], skip_special_tokens=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()
         
     except Exception as e:
-        return f"エラー: {str(e)}"
+        return f"エラーが発生しました: {str(e)}"
+
+demo = gr.ChatInterface(
+    fn=predict,
+    title="💬 Gemma 2 for ELYZA-tasks",
+    description="ELYZA-tasks-100-TV用に最適化された日本語LLMです"
+)
 
-demo = gr.ChatInterface(fn=predict)
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=True)