Update app.py

app.py

@@ -60,43 +60,44 @@ def load_settings_from_ini(filename='character_settings.ini'):
 
 # LlamaCppのラッパークラス
 class LlamaCppAdapter:
-    @spaces.GPU(duration=120)    
     def __init__(self, model_path, n_ctx=4096):
-        print(f"モデルの初期化: {model_path}")
+        self.model_path = model_path
+        self.n_ctx = n_ctx
         self.llama = None
-        try:
-            self.llama = Llama(model_path=model_path, n_ctx=n_ctx, n_gpu_layers=0)
-        except Exception as e:
-            print(f"モデルの初期化中にエラーが発生しました: {e}")
 
-    def generate(self, prompt, max_new_tokens=4096, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
+    @spaces.GPU(duration=120)
+    def initialize_and_generate(self, prompt, max_new_tokens=4096, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
         if self.llama is None:
-            return {"choices": [{"text": "モデルの初期化に失敗しました。"}]}
+            print(f"モデルの初期化: {self.model_path}")
+            try:
+                self.llama = Llama(model_path=self.model_path, n_ctx=self.n_ctx, n_gpu_layers=-1)
+                print("モデルの初期化が成功しました。")
+            except Exception as e:
+                print(f"モデルの初期化中にエラーが発生しました: {e}")
+                return {"choices": [{"text": "モデルの初期化に失敗しました。"}]}
+
         try:
-            return self._generate(prompt, max_new_tokens, temperature, top_p, top_k, stop)
+            result = self.llama(
+                prompt,
+                temperature=temperature,
+                max_tokens=max_new_tokens,
+                top_p=top_p,
+                top_k=top_k,
+                stop=stop,
+                repeat_penalty=1.2,
+            )
+            return result
         except Exception as e:
             print(f"生成中にエラーが発生しました: {e}")
             return {"choices": [{"text": "生成中にエラーが発生しました。"}]}
 
-    def _generate(self, prompt: str, max_new_tokens: int, temperature: float, top_p: float, top_k: int, stop: list):
-        return self.llama(
-            prompt,
-            temperature=temperature,
-            max_tokens=max_new_tokens,
-            top_p=top_p,
-            top_k=top_k,
-            stop=stop,
-            repeat_penalty=1.2,
-        )
+    def generate(self, prompt, max_new_tokens=4096, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
+        return self.initialize_and_generate(prompt, max_new_tokens, temperature, top_p, top_k, stop)
 
 # キャラクターメーカークラス
 class CharacterMaker:
     def __init__(self):
-        self.llama = None
-        try:
-            self.llama = LlamaCppAdapter(model_path)
-        except Exception as e:
-            print(f"LlamaCppAdapter の初期化中にエラーが発生しました: {e}")
+        self.llama = LlamaCppAdapter(model_path)
         self.history = []
         self.settings = load_settings_from_ini()
         if not self.settings:
@@ -129,15 +130,15 @@ class CharacterMaker:
             save_settings_to_ini(self.settings)
 
     def make(self, input_str: str):
-        if self.llama is None:
-            return "申し訳ありません。モデルが正しく初期化されていません。"
-        
         prompt = self._generate_aki(input_str)
         print(prompt)
         print("-----------------")
         try:
             res = self.llama.generate(prompt, max_new_tokens=1000, stop=["<END>", "\n"])
-            res_text = res["choices"][0]["text"]
+            if isinstance(res, dict) and "choices" in res and len(res["choices"]) > 0:
+                res_text = res["choices"][0]["text"]
+            else:
+                res_text = "応答の生成に失敗しました。"
         except Exception as e:
             print(f"生成中にエラーが発生しました: {e}")
             res_text = "申し訳ありません。応答の生成中にエラーが発生しました。"
@@ -170,11 +171,7 @@ class CharacterMaker:
 
     def reset(self):
         self.history = []
-        try:
-            self.llama = LlamaCppAdapter(model_path)
-        except Exception as e:
-            print(f"LlamaCppAdapter の再初期化中にエラーが発生しました: {e}")
-            self.llama = None
+        self.llama = LlamaCppAdapter(model_path)
 
 character_maker = CharacterMaker()