app.py

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)