app.py

from flask import Flask, request, Response
import logging
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
import sys

llm = LlmBackend()
_lock = threading.Lock()
    
SYSTEM_PROMPT = os.environ.get('SYSTEM_PROMPT') or "Ты — русскоязычный автоматический ассистент. Ты максимально точно и отвечаешь на запросы пользователя, используя русский язык."
CONTEXT_SIZE = int(os.environ.get('CONTEXT_SIZE', '500'))
HF_CACHE_DIR = os.environ.get('HF_CACHE_DIR') or '/home/user/app/.cache'
USE_SYSTEM_PROMPT = os.environ.get('USE_SYSTEM_PROMPT', '').lower() == "true" or False
ENABLE_GPU = os.environ.get('ENABLE_GPU', '').lower() == "true" or False
GPU_LAYERS = int(os.environ.get('GPU_LAYERS', '0'))
CHAT_FORMAT = os.environ.get('CHAT_FORMAT') or 'llama-2'
REPO_NAME = os.environ.get('REPO_NAME') or 'IlyaGusev/saiga2_7b_gguf'
MODEL_NAME = os.environ.get('MODEL_NAME') or 'model-q4_K.gguf'
DATASET_REPO_URL = os.environ.get('DATASET_REPO_URL') or "https://huggingface.co/datasets/muryshev/saiga-chat"
DATA_FILENAME = os.environ.get('DATA_FILENAME') or "data-saiga-cuda-release.xml"
HF_TOKEN = os.environ.get("HF_TOKEN")

# Create a lock object
lock = threading.Lock()
app = Flask('llm_api')

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
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"

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, cache_dir=HF_CACHE_DIR) + '/' + MODEL_NAME
app.logger.info('Model path: ' + MODEL_PATH)

DATA_FILE = os.path.join("dataset", DATA_FILENAME)


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)

@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 called')
    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, 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)

def init_model():
    llm.load_model(model_path=MODEL_PATH, context_size=CONTEXT_SIZE, enable_gpu=ENABLE_GPU, gpu_layer_number=GPU_LAYERS)
    
# 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
        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__":
    
    init_model()
    
    # scheduler = BackgroundScheduler()
    # scheduler.add_job(check_last_request_time, trigger='interval', minutes=1)
    # scheduler.start()
    
    app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)