Update app.py

app.py

@@ -3,6 +3,8 @@ import logging
 from llama_cpp import Llama
 import threading
 from huggingface_hub import snapshot_download
+import gc
+import os.path
 
 SYSTEM_PROMPT = "Ты — русскоязычный автоматический ассистент. Ты максимально точно и отвечаешь на запросы пользователя, используя русский язык."
 SYSTEM_TOKEN = 1788
@@ -17,6 +19,7 @@ ROLE_TOKENS = {
 }
 
 CONTEXT_SIZE = 4000
+ENABLE_GPU = False
 
 # Create a lock object
 lock = threading.Lock()
@@ -35,22 +38,52 @@ app.logger.setLevel(logging.DEBUG)  # Set the desired logging level
 repo_name = "IlyaGusev/saiga2_70b_gguf"
 model_name = "ggml-model-q4_1.gguf"
 
+#repo_name = "IlyaGusev/saiga2_7b_gguf"
+#model_name = "model-q4_K.gguf"
+local_dir = '.'
+
+if os.path.isdir('/data'):
+    app.logger.info('Persistent storage enabled')
 
+model = None
 
-snapshot_download(repo_id=repo_name, local_dir=".", allow_patterns=model_name)
+model_path = snapshot_download(repo_id=repo_name, allow_patterns=model_name) + '/' + model_name
+app.logger.info('Model path: ' + model_path)
 
-model = Llama(
-    model_path=model_name,
-    n_ctx=CONTEXT_SIZE,
-    n_parts=1,
-    #n_batch=100,
-    logits_all=True,
-    n_threads=12,
-    verbose=True,
-    n_gpu_layers=35,
-    n_gqa=8       #must be set for 70b models
-)
+def init_model(context_size, enable_gpu=False, gpu_layer_number=35):
+    global model
+    
+    if model is not None:
+        del model
+        gc.collect()
+    
+    if enable_gpu:
+        model = Llama(
+            model_path=model_path,
+            n_ctx=context_size,
+            n_parts=1,
+            #n_batch=100,
+            logits_all=True,
+            #n_threads=12,
+            verbose=True,
+            n_gpu_layers=gpu_layer_number,
+            n_gqa=8       #must be set for 70b models
+        )
+        return model
+    else:
+        model = Llama(
+            model_path=model_path,
+            n_ctx=context_size,
+            n_parts=1,
+            #n_batch=100,
+            logits_all=True,
+            #n_threads=12,
+            verbose=True,
+            n_gqa=8       #must be set for 70b models
+        )
+        return model
 
+init_model(CONTEXT_SIZE, ENABLE_GPU, 35)
 
 def get_message_tokens(model, role, content):
     message_tokens = model.tokenize(content.encode("utf-8"))
@@ -73,10 +106,10 @@ def get_system_tokens_for_preprompt(model, preprompt):
     }
     return get_message_tokens(model, **system_message)
 
-app.logger.info('Evaluating system tokens start')
+#app.logger.info('Evaluating system tokens start')
 #system_tokens = get_system_tokens(model)
 #model.eval(system_tokens)
-app.logger.info('Evaluating system tokens end')
+#app.logger.info('Evaluating system tokens end')
 
 stop_generation = False
 
@@ -96,8 +129,19 @@ def generate_tokens(model, generator):
                 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
+    stop_generation = True
+
+    new_size = int(request.args.get('size', CONTEXT_SIZE))
+    init_model(new_size, enable_gpu=ENABLE_GPU)
+    
+    return Response('Size changed', content_type='text/plain')   
+    
 @app.route('/stop_generation', methods=['GET'])
 def handler_stop_generation():
     global stop_generation