Another bunch of fixes

app.py

@@ -11,12 +11,16 @@ from apscheduler.schedulers.background import BackgroundScheduler
 from datetime import datetime, timedelta
 from llm_backend import LlmBackend
 import json
+import log
+import sys
 
 llm = LlmBackend()
 _lock = threading.Lock()
     
 SYSTEM_PROMPT = os.environ.get('SYSTEM_PROMPT') or "Ты — русскоязычный автоматический ассистент. Ты максимально точно и отвечаешь на запросы пользователя, используя русский язык."
-CONTEXT_SIZE = os.environ.get('CONTEXT_SIZE') or 500 
+CONTEXT_SIZE = os.environ.get('CONTEXT_SIZE') or 500
+HF_CACHE_DIR = os.environ.get('HF_CACHE_DIR') or '/root/.cache'
+USE_SYSTEM_PROMPT = os.environ.get('USE_SYSTEM_PROMPT') or False
 ENABLE_GPU = os.environ.get('ENABLE_GPU') or False
 GPU_LAYERS = os.environ.get('GPU_LAYERS') or 0
 N_GQA = os.environ.get('N_GQA') or None #must be set to 8 for 70b models
@@ -24,9 +28,12 @@ CHAT_FORMAT = os.environ.get('CHAT_FORMAT') or 'llama-2'
 
 # Create a lock object
 lock = threading.Lock()
+app = Flask('llm_api')
 
-app = Flask(__name__)
-# Configure Flask logging
+app.logger.handlers.clear()
+handler = logging.StreamHandler(sys.stdout)
+handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+app.logger.addHandler(handler)
 app.logger.setLevel(logging.DEBUG)
 
 # Variable to store the last request time
@@ -51,7 +58,7 @@ if os.path.isdir('/data'):
 
 model = None
 
-MODEL_PATH = snapshot_download(repo_id=repo_name, allow_patterns=model_name) + '/' + model_name
+MODEL_PATH = snapshot_download(repo_id=repo_name, allow_patterns=model_name, cache_dir=HF_CACHE_DIR) + '/' + model_name
 app.logger.info('Model path: ' + MODEL_PATH)
 
 DATASET_REPO_URL = "https://huggingface.co/datasets/muryshev/saiga-chat"
@@ -81,25 +88,6 @@ app.logger.info("hfh: "+huggingface_hub.__version__)
 #         commit_url = repo.push_to_hub()
 #         app.logger.info(commit_url)
 
-def generate_tokens(model, generator):
-    global stop_generation
-    app.logger.info('generate_tokens started')
-    with lock:
-        try:
-            for token in generator:            
-                if token == model.token_eos() or stop_generation:
-                    stop_generation = False
-                    app.logger.info('End generating')
-                    yield b''  # End of chunk
-                    break
-                    
-                token_str = model.detokenize([token])#.decode("utf-8", errors="ignore")
-                yield token_str 
-        except Exception as e:
-            app.logger.info('generator exception')
-            app.logger.info(e)
-            yield b''  # End of chunk
-
 @app.route('/change_context_size', methods=['GET'])
 def handler_change_context_size():
     global stop_generation, model
@@ -142,12 +130,7 @@ def generate_and_log_tokens(user_request, generator):
 @app.route('/', methods=['POST'])
 def generate_response():
 
-    app.logger.info('generate_response')
-    with _lock:
-        if not llm.is_model_loaded():
-            app.logger.info('model loading')    
-            init_model()
-        
+    app.logger.info('generate_response called')
     data = request.get_json()
     app.logger.info(data)
     messages = data.get("messages", [])
@@ -165,12 +148,9 @@ def generate_response():
         'return_full_text': parameters.get("return_full_text", False)
     }
     
-    generator = llm.create_chat_generator_for_saiga(messages=messages, parameters=p)
+    generator = llm.create_chat_generator_for_saiga(messages=messages, parameters=p, use_system_prompt=USE_SYSTEM_PROMPT)
     app.logger.info('Generator created')
 
-
-    
-
     # Use Response to stream tokens
     return Response(generate_and_log_tokens(user_request='1', generator=generator), content_type='text/plain', status=200, direct_passthrough=True)
 
@@ -182,7 +162,6 @@ def check_last_request_time():
     global last_request_time
     current_time = datetime.now()
     if (current_time - last_request_time).total_seconds() > 300:  # 5 minutes in seconds
-        # Perform the action (e.g., set a variable)
         llm.unload_model()
         app.logger.info(f"Model unloaded at {current_time}")
     else:
@@ -190,10 +169,11 @@ def check_last_request_time():
 
 
 if __name__ == "__main__":
-    scheduler = BackgroundScheduler()
-    scheduler.add_job(check_last_request_time, trigger='interval', minutes=1)
-    scheduler.start()
     
     init_model()
     
-    app.run(host="0.0.0.0", port=7860, debug=True, threaded=True)
+    app.run(host="0.0.0.0", port=7860, debug=True, threaded=True)
+    
+    scheduler = BackgroundScheduler()
+    scheduler.add_job(check_last_request_time, trigger='interval', minutes=1)
+    scheduler.start()

llm_backend.py

@@ -1,7 +1,11 @@
 from llama_cpp import Llama
 import gc
 import threading
+import logging
+import sys
 
+log = logging.getLogger('llm_api.backend')
+    
 class LlmBackend:
     
     SYSTEM_PROMPT = "Ты — русскоязычный автоматический ассистент. Ты максимально точно и отвечаешь на запросы пользователя, используя русский язык."
@@ -18,6 +22,7 @@ class LlmBackend:
 
     _instance = None
     _model = None
+    _model_params = None
     _lock = threading.Lock()
     
     def __new__(cls):
@@ -30,6 +35,14 @@ class LlmBackend:
         return self._model is not None
     
     def load_model(self, model_path, context_size=2000, enable_gpu=True, gpu_layer_number=35, n_gqa=8, chat_format='llama-2'):
+        log.info('load_model - started')
+        self._model_params = {}
+        self._model_params['model_path'] = model_path
+        self._model_params['context_size'] = context_size
+        self._model_params['enable_gpu'] = enable_gpu
+        self._model_params['gpu_layer_number'] = gpu_layer_number
+        self._model_params['n_gqa'] = n_gqa
+        self._model_params['chat_format'] = chat_format
         
         if self._model is not None:
             self.unload_model()
@@ -44,10 +57,11 @@ class LlmBackend:
                     #n_batch=100,
                     logits_all=True,
                     #n_threads=12,
-                    verbose=True,
+                    verbose=False,
                     n_gpu_layers=gpu_layer_number,
                     n_gqa=n_gqa       #must be set for 70b models
                 )
+                log.info('load_model - finished')
                 return self._model
             else:
                 self._model = Llama(
@@ -58,9 +72,10 @@ class LlmBackend:
                     #n_batch=100,
                     logits_all=True,
                     #n_threads=12,
-                    verbose=True,
+                    verbose=False,
                     n_gqa=n_gqa       #must be set for 70b models
                 )
+                log.info('load_model - finished')
                 return self._model
         
     def set_system_prompt(self, prompt):
@@ -68,54 +83,71 @@ class LlmBackend:
             self.SYSTEM_PROMPT = prompt
         
     def unload_model(self):
+        log.info('unload_model - started')
         with self._lock:
             if self._model is not None:
                 del self._model
-            
+        log.info('unload_model - finished')
+    
+    def ensure_model_is_loaded(self):
+        log.info('ensure_model_is_loaded - started')
+        if not self.is_model_loaded():
+            log.info('ensure_model_is_loaded - model reloading')
+            if self._model_params is not None:
+                self.load_model(**self._model_params)
+            else:
+                log.info('ensure_model_is_loaded - No model config found. Reloading can not be done.')
+        log.info('ensure_model_is_loaded - finished')
+                    
     def generate_tokens(self, generator):
-        print('generate_tokens called')
+        log.info('generate_tokens - started')
         with self._lock:
-            print('generate_tokens started')
+            self.ensure_model_is_loaded()
+                
             try:
                 for token in generator:            
                     if token == self._model.token_eos():
-                        print('End generating')
+                        log.info('generate_tokens - finished')
                         yield b''  # End of chunk
                         break
                         
                     token_str = self._model.detokenize([token])#.decode("utf-8", errors="ignore")
                     yield token_str 
             except Exception as e:
-                print('generator exception')
-                print(e)
+                log.error('generate_tokens - error')
+                log.error(e)
                 yield b''  # End of chunk
                 
     def create_chat_completion(self, messages, stream=True):
-        print('create_chat_completion called')
+        log.info('create_chat_completion called')
         with self._lock:
-            print('create_chat_completion started')
+            log.info('create_chat_completion started')
             try:
                 return self._model.create_chat_completion(messages=messages, stream=stream)
             except Exception as e:
-                print('create_chat_completion exception')
-                print(e)
+                log.error('create_chat_completion - error')
+                log.error(e)
                 return None
                 
     
     def get_message_tokens(self, role, content):
+        log.info('get_message_tokens - started')
+        self.ensure_model_is_loaded()
         message_tokens = self._model.tokenize(content.encode("utf-8"))
         message_tokens.insert(1, self.ROLE_TOKENS[role])
         message_tokens.insert(2, self.LINEBREAK_TOKEN)
         message_tokens.append(self._model.token_eos())
+        log.info('get_message_tokens - finished')
         return message_tokens
 
     def get_system_tokens(self):
         return self.get_message_tokens(role="system", content=self.SYSTEM_PROMPT)
     
-    def create_chat_generator_for_saiga(self, messages, parameters):
-        print('create_chat_completion called')
+    def create_chat_generator_for_saiga(self, messages, parameters, use_system_prompt=True):
+        log.info('create_chat_generator_for_saiga - started')
         with self._lock:
-            tokens = self.get_system_tokens()
+            self.ensure_model_is_loaded()
+            tokens = self.get_system_tokens() if use_system_prompt else []
             for message in messages:
                 message_tokens = self.get_message_tokens(role=message.get("from"), content=message.get("content", ""))
                 tokens.extend(message_tokens)
@@ -128,19 +160,25 @@ class LlmBackend:
                 temp=parameters['temperature'],
                 repeat_penalty=parameters['repetition_penalty']
             )
+            log.info('create_chat_generator_for_saiga - finished')
             return generator
         
     def generate_tokens(self, generator):
+        log.info('generate_tokens - started')
         with self._lock:
+            self.ensure_model_is_loaded()
             try:
                 for token in generator:            
                     if token == self._model.token_eos():
                         yield b''  # End of chunk
+                        log.info('generate_tokens - finished')
                         break
                         
                     token_str = self._model.detokenize([token])#.decode("utf-8", errors="ignore")
                     yield token_str 
             except Exception as e:
+                log.error('generate_tokens - error')
+                log.error(e)
                 yield b''  # End of chunk