init

.gitignore

@@ -0,0 +1,3 @@
+.env/
+data/
+.vscode/

Dockerfile

@@ -0,0 +1,32 @@
+ARG CUDA_IMAGE="12.1.1-devel-ubuntu22.04"
+FROM nvidia/cuda:${CUDA_IMAGE}
+
+# We need to set the host to 0.0.0.0 to allow outside access
+ENV HOST 0.0.0.0
+RUN useradd -m -u 1000 user
+WORKDIR /home/user/app
+COPY --link --chown=1000 ./ /home/user/app
+
+RUN apt-get update && apt-get upgrade -y \
+    && apt-get install -y git git-lfs build-essential \
+    python3 python3-pip gcc wget \
+    ocl-icd-opencl-dev opencl-headers clinfo \
+    libclblast-dev libopenblas-dev \
+    && mkdir -p /etc/OpenCL/vendors && echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd
+
+# setting build related env vars
+ENV CUDA_DOCKER_ARCH=all
+ENV LLAMA_CUBLAS=1
+
+# Install depencencies
+RUN python3 -m pip install --no-cache-dir --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings starlette-context huggingface-hub==0.14.1 flask apscheduler
+
+# Install llama-cpp-python (build with cuda)
+RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install --no-cache-dir llama-cpp-python
+
+RUN git config --global user.email "amatenkov@ntr.ai"
+RUN git config --global user.name "Andrew Matenkov"
+
+EXPOSE 7860
+# Run the server
+CMD python3 -m app

app.py

@@ -0,0 +1,200 @@
+from flask import Flask, request, Response
+import logging
+from llama_cpp import Llama
+import threading
+from huggingface_hub import snapshot_download#, Repository
+import huggingface_hub
+import gc
+import os.path
+import xml.etree.ElementTree as ET
+from apscheduler.schedulers.background import BackgroundScheduler
+from datetime import datetime, timedelta
+from llm_backend import LlmBackend
+import json
+
+llm = LlmBackend()
+_lock = threading.Lock()
+    
+SYSTEM_PROMPT = os.environ.get('SYSTEM_PROMPT') or "Ты — русскоязычный автоматический ассистент. Ты максимально точно и отвечаешь на запросы пользователя, используя русский язык."
+
+CONTEXT_SIZE = os.environ.get('CONTEXT_SIZE') or 500 
+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
+CHAT_FORMAT = os.environ.get('CHAT_FORMAT') or 'llama-2'
+
+# Create a lock object
+lock = threading.Lock()
+
+app = Flask(__name__)
+# Configure Flask logging
+app.logger.setLevel(logging.DEBUG)
+
+# Variable to store the last request time
+last_request_time = datetime.now()
+
+# Initialize the model when the application starts
+#model_path = "../models/model-q4_K.gguf"  # Replace with the actual model path
+#model_name = "model/ggml-model-q4_K.gguf"
+
+#repo_name = "IlyaGusev/saiga2_13b_gguf"
+#model_name = "model-q4_K.gguf"
+
+#epo_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
+
+MODEL_PATH = snapshot_download(repo_id=repo_name, allow_patterns=model_name) + '/' + model_name
+app.logger.info('Model path: ' + MODEL_PATH)
+
+DATASET_REPO_URL = "https://huggingface.co/datasets/muryshev/saiga-chat"
+DATA_FILENAME = "data-saiga-cuda-release.xml"
+DATA_FILE = os.path.join("dataset", DATA_FILENAME)
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+app.logger.info("hfh: "+huggingface_hub.__version__)
+
+# repo = Repository(
+#     local_dir="dataset", clone_from=DATASET_REPO_URL, use_auth_token=HF_TOKEN
+# )
+
+
+
+# def log(req: str = '', resp: str = ''):
+#     if req or resp:
+#         element = ET.Element("row", {"time": str(datetime.now()) })
+#         req_element = ET.SubElement(element, "request")
+#         req_element.text = req
+#         resp_element = ET.SubElement(element, "response")
+#         resp_element.text = resp
+    
+#         with open(DATA_FILE, "ab+") as xml_file:
+#             xml_file.write(ET.tostring(element, encoding="utf-8"))
+        
+#         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
+    stop_generation = True
+
+    new_size = int(request.args.get('size', CONTEXT_SIZE))
+    init_model(new_size, ENABLE_GPU, GPU_LAYERS)
+    
+    return Response('Size changed', content_type='text/plain')   
+    
+@app.route('/stop_generation', methods=['GET'])
+def handler_stop_generation():
+    global stop_generation
+    stop_generation = True
+    return Response('Stopped', content_type='text/plain')        
+                
+@app.route('/', methods=['GET', 'PUT', 'DELETE', 'PATCH'])
+def generate_unknown_response():
+    app.logger.info('unknown method: '+request.method)
+    try:
+        request_payload = request.get_json()
+        app.logger.info('payload: '+request.get_json())
+    except Exception as e:
+        app.logger.info('payload empty')
+
+    return Response('What do you want?', content_type='text/plain')
+
+response_tokens = bytearray()
+def generate_and_log_tokens(user_request, generator):
+    global response_tokens, last_request_time
+    for token in llm.generate_tokens(generator):
+        if token == b'': # or (max_new_tokens is not None and i >= max_new_tokens):
+            last_request_time = datetime.now()
+            # log(json.dumps(user_request), response_tokens.decode("utf-8", errors="ignore"))
+            response_tokens = bytearray()
+            break
+        response_tokens.extend(token)
+        yield token
+            
+@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()
+        
+    data = request.get_json()
+    app.logger.info(data)
+    messages = data.get("messages", [])
+    preprompt = data.get("preprompt", "")
+    parameters = data.get("parameters", {})
+    
+    # Extract parameters from the request
+    p = {
+        'temperature': parameters.get("temperature", 0.01),
+        'truncate': parameters.get("truncate", 1000),
+        'max_new_tokens': parameters.get("max_new_tokens", 1024),
+        'top_p': parameters.get("top_p", 0.85),
+        'repetition_penalty': parameters.get("repetition_penalty", 1.2),
+        'top_k': parameters.get("top_k", 30),
+        'return_full_text': parameters.get("return_full_text", False)
+    }
+    
+    generator = llm.create_chat_generator_for_saiga(messages=messages, parameters=p)
+    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)
+
+def init_model():
+    llm.load_model(model_path=MODEL_PATH, context_size=CONTEXT_SIZE, enable_gpu=ENABLE_GPU, gpu_layer_number=GPU_LAYERS, n_gqa=N_GQA)
+    
+# Function to check if no requests were made in the last 5 minutes
+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:
+        app.logger.info(f"No action needed at {current_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)

llm_backend.py

@@ -0,0 +1,152 @@
+from llama_cpp import Llama
+import gc
+import threading
+
+class LlmBackend:
+    
+    SYSTEM_PROMPT = "Ты — русскоязычный автоматический ассистент. Ты максимально точно и отвечаешь на запросы пользователя, используя русский язык."
+    SYSTEM_TOKEN = 1788
+    USER_TOKEN = 1404
+    BOT_TOKEN = 9225
+    LINEBREAK_TOKEN = 13
+
+    ROLE_TOKENS = {
+        "user": USER_TOKEN,
+        "bot": BOT_TOKEN,
+        "system": SYSTEM_TOKEN
+    }
+
+    _instance = None
+    _model = None
+    _lock = threading.Lock()
+    
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(LlmBackend, cls).__new__(cls)
+        return cls._instance
+    
+    
+    def is_model_loaded(self):
+        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'):
+        
+        if self._model is not None:
+            self.unload_model()
+            
+        with self._lock:    
+            if enable_gpu:
+                self._model = Llama(
+                    model_path=model_path,
+                    chat_format=chat_format,
+                    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=n_gqa       #must be set for 70b models
+                )
+                return self._model
+            else:
+                self._model = Llama(
+                    model_path=model_path,
+                    chat_format=chat_format,
+                    n_ctx=context_size,
+                    n_parts=1,
+                    #n_batch=100,
+                    logits_all=True,
+                    #n_threads=12,
+                    verbose=True,
+                    n_gqa=n_gqa       #must be set for 70b models
+                )
+                return self._model
+        
+    def set_system_prompt(self, prompt):
+        with self._lock:
+            self.SYSTEM_PROMPT = prompt
+        
+    def unload_model(self):
+        with self._lock:
+            if self._model is not None:
+                self._model.llama_free_model()
+                del self._model
+            
+    def generate_tokens(self, generator):
+        print('generate_tokens called')
+        with self._lock:
+            print('generate_tokens started')
+            try:
+                for token in generator:            
+                    if token == self._model.token_eos():
+                        print('End generating')
+                        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)
+                yield b''  # End of chunk
+                
+    def create_chat_completion(self, messages, stream=True):
+        print('create_chat_completion called')
+        with self._lock:
+            print('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)
+                return None
+                
+    
+    def get_message_tokens(self, role, content):
+        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())
+        return message_tokens
+
+    def get_system_tokens(self):
+        system_message = {
+            "role": "system",
+            "content": self.SYSTEM_PROMPT
+        }
+        return self.get_message_tokens(self._model, **system_message)
+    
+    def create_chat_generator_for_saiga(self, messages, parameters):
+        print('create_chat_completion called')
+        with self._lock:
+            tokens = self.get_system_tokens()
+            for message in messages:
+                message_tokens = self.get_message_tokens(role=message.get("from"), content=message.get("content", ""))
+                tokens.extend(message_tokens)
+            
+            tokens.extend([self._model.token_bos(), self.BOT_TOKEN, self.LINEBREAK_TOKEN])
+            generator = self._model.generate(
+                tokens,
+                top_k=parameters['top_k'],
+                top_p=parameters['top_p'],
+                temp=parameters['temperature'],
+                repeat_penalty=parameters['repetition_penalty']
+            )
+            return generator
+        
+    def generate_tokens(self, generator):
+        with self._lock:
+            try:
+                for token in generator:            
+                    if token == self._model.token_eos():
+                        yield b''  # End of chunk
+                        break
+                        
+                    token_str = self._model.detokenize([token])#.decode("utf-8", errors="ignore")
+                    yield token_str 
+            except Exception as e:
+                yield b''  # End of chunk
+
+    
+