Update app.py

app.py

@@ -6,9 +6,7 @@ from llama_cpp import Llama
 import configparser
 from functools import partial
 from utils.dl_utils import dl_guff_model
-import io
-import tempfile
-import csv
+import threading
 
 # 定数
 DEFAULT_INI_FILE = 'settings.ini'
@@ -34,7 +32,6 @@ model_path = os.path.join("models", model_filename)
 if not os.path.exists(model_path):
     dl_guff_model("models", f"https://huggingface.co/MCZK/EZO-Common-9B-gemma-2-it-GGUF/resolve/main/{model_filename}")
 
-
 class ConfigManager:
     @staticmethod
     def load_settings(filename):
@@ -106,6 +103,9 @@ class TextGenerator:
         self.llm = None
         self.settings = None
         self.current_model = None
+        self.history = []
+        self.use_chat_format = False
+        self.model_loaded = threading.Event()
 
     @spaces.GPU(duration=120)
     def load_model(self):
@@ -113,50 +113,79 @@ class TextGenerator:
             del self.llm
             self.llm = None
 
+        self.model_loaded.clear()
         try:
             model_path = os.path.join(MODEL_DIR, self.settings['DEFAULT_GEN_MODEL'])
             n_gpu_layers = self.settings['gen_n_gpu_layers']
             self.llm = Llama(model_path=model_path, n_ctx=self.settings['gen_n_ctx'], n_gpu_layers=n_gpu_layers)
             self.current_model = 'GEN'
             print(f"GEN モデル {model_path} のロードが完了しました。(n_gpu_layers: {n_gpu_layers})")
+            self.model_loaded.set()
         except Exception as e:
             print(f"GEN モデルのロード中にエラーが発生しました: {str(e)}")
+            self.model_loaded.set()
 
     def generate_text(self, text, gen_characters, gen_token_multiplier, instruction):
         if not self.llm:
             self.load_model()
         
-        if not self.llm:
+        if not self.model_loaded.wait(timeout=30) or not self.llm:
+            return "モデルのロードに失敗しました。設定を確認してください。"
+        
+        input_str = f"{instruction}\n\n生成するテキスト（目安は{gen_characters}文字）：\n\n{text}"
+        return self.generate_response(input_str, gen_characters, gen_token_multiplier)
+
+    def generate_response(self, input_str, gen_characters, gen_token_multiplier):
+        if not self.model_loaded.wait(timeout=30) or not self.llm:
             return "モデルのロードに失敗しました。設定を確認してください。"
         
-        author_description = self.settings.get('gen_author_description', '')
         max_tokens = int(gen_characters * gen_token_multiplier)
         
-        messages = [
-            {"role": "user", "content": f"{author_description}\n\n以下の指示に従ってテキストを生成してください："},
-            {"role": "assistant", "content": "はい、承知しました。指示に従ってテキストを生成いたします。"},
-            {"role": "user", "content": f"{instruction}\n\n生成するテキスト（目安は{gen_characters}文字）：\n\n{text}"}
-        ]
+        if not self.use_chat_format:
+            try:
+                prompt = self._generate_prompt(input_str)
+                res = self.llm.create_completion(prompt=prompt, max_tokens=max_tokens)
+                res_text = res["choices"][0]["text"]
+                self.history.append({"user": input_str, "assistant": res_text})
+                return res_text
+            except Exception as e:
+                print(f"既存の形式でのレスポンス生成に失敗しました: {str(e)}")
+                print("チャット形式に切り替えます。")
+                self.use_chat_format = True
         
-        try:
-            response = self.llm.create_chat_completion(
-                messages=messages,
-                max_tokens=max_tokens,
-                temperature=self.settings['gen_temperature'],
-                top_p=self.settings['gen_top_p'],
-                top_k=self.settings['gen_top_k'],
-                repeat_penalty=self.settings['gen_rep_pen'],
-            )
+        if self.use_chat_format:
+            chat_messages = [
+                {"role": "system", "content": self.settings.get('gen_author_description', '')},
+                {"role": "user", "content": input_str}
+            ]
             
-            generated_text = response["choices"][0]["message"]["content"].strip()
-            return generated_text
-        except Exception as e:
-            print(f"テキスト生成中にエラーが発生しました: {str(e)}")
-            return "テキスト生成中にエラーが発生しました。設定を確認してください。"
+            try:
+                response = self.llm.create_chat_completion(
+                    messages=chat_messages,
+                    max_tokens=max_tokens,
+                    temperature=self.settings['gen_temperature'],
+                    top_p=self.settings['gen_top_p'],
+                    top_k=self.settings['gen_top_k'],
+                    repeat_penalty=self.settings['gen_rep_pen'],
+                )
+                
+                res_text = response["choices"][0]["message"]["content"].strip()
+                self.history.append({"user": input_str, "assistant": res_text})
+                return res_text
+            except Exception as e:
+                print(f"チャット形式でのレスポンス生成に失敗しました: {str(e)}")
+                return "レスポンス生成中にエラーが発生しました。設定を確認してください。"
+
+    def _generate_prompt(self, input_str):
+        return f"{self.settings.get('gen_author_description', '')}\n\nUser: {input_str}\nAssistant:"
 
     def load_settings(self, filename):
         self.settings = Settings.load_from_ini(filename)
 
+    def reset(self):
+        self.history = []
+        self.use_chat_format = False
+
 # グローバル変数
 text_generator = TextGenerator()
 model_files = ModelManager.get_model_files()
@@ -273,6 +302,8 @@ def build_gradio_interface():
                         input_component = gr.Textbox(label=key, value=value)
                     
                     input_component.change(
+                        partial(ConfigManager.update_setting, 'GenerateParameters', key),
+                        input_component.change(
                         partial(ConfigManager.update_setting, 'GenerateParameters', key),
                         inputs=[input_component],
                         outputs=[output]
@@ -294,4 +325,4 @@ if __name__ == "__main__":
     text_generator.load_settings(DEFAULT_INI_FILE)
     
     demo = build_gradio_interface()
-    demo.launch(share=True)
+    demo.launch(share=True)