#https://huggingface.co/Qwen/Qwen3-4B
#https://github.com/QwenLM/Qwen3

import torch
import os

from transformers import AutoModelForCausalLM, AutoTokenizer

import time 
import math
import requests
import json
import argparse
import logging

from flask import Flask, abort, jsonify, request

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,
)
logger = logging.getLogger(__name__)

# Run on CPU
runOnCpu = False
if runOnCpu:
    device = "cpu"
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
else:
    device = "cuda"

def clean(answer):
    i = answer.find("[/INST]")
    if i > 0:
        answer = answer[i + 8:]
    i = answer.find("</s>")
    if i > 0:
        answer = answer[0:i]
    answer = str(answer)    
    answer = answer.replace('\\n', '\n').replace('\\t', '\t').replace('"', '')
    return answer
             
def run(data):
    if 'task' not in data or len(data['task']) == 0: abort(500)
    task = data['task']
    if 'max_tokens' in data: 
        max_tokens = data['max_tokens']
    elif 'max_new_tokens' in data: 
        max_tokens = data['max_new_tokens']
    else:
        max_tokens = 100
    if 'min_tokens' in data: 
        min_tokens = data['min_tokens']
    else:
        min_tokens = 1
    if 'num_beams' in data: 
        num_beams = data['num_beams']
    else:
        num_beams=2
    if 'early_stopping' in data: 
        early_stopping = data['early_stopping']
    else:
        early_stopping = True
    if 'max_time' in data: 
        max_time = data['max_time']
    else:
        max_time = 10
    if 'temperature' in data: 
        temperature = data['temperature']
    else:
        temperature = 0.15
    if 'top_p' in data: 
        top_p = data['top_p']
    else:
        top_p = 0.5
    if 'enable_thinking' in data: 
        enable_thinking = data['enable_thinking']
    else:
        enable_thinking = False
    
    if 'n' in data: 
        num_array = data['n']
#        sampling_params = SamplingParams(
#            temperature=temperature,
#            top_p=top_p,
#            min_tokens=min_tokens,
#            max_tokens=max_tokens,
#            top_k=20,
#            repetition_penalty=1.5,
#            n = num_array)
    else:
        num_array = 0 
#        sampling_params = SamplingParams(
#            temperature=temperature,
#            top_p=top_p,
#            min_tokens=min_tokens,
#            max_tokens=max_tokens)
    
    messages = [
#        {"role": "system", "content": "You are a helpful assistant answering only in English or Russian languages."},
        {"role": "user", "content": task}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # conduct text completion
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=32768
    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
    generated_text = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")

#    if num_array == 0:
#        generated_text = clean(result.outputs[0].text)
#        torch.cuda.empty_cache()
#    else:
#        generated_text = []
#        for output in result.outputs:
#            generated_text.append(clean(output.text))        
#    torch.cuda.empty_cache()

    logger.info(f"Generated text: {generated_text!r}")
    return generated_text


#parser = argparse.ArgumentParser()
#parser.add_argument('--max_model_len', type=int, default=8192, help='Max model length') #8192
#parser.add_argument('--gpu_memory_utilization', type=float, default=0.9, help='GPU memory utilization') #0.9
#args = parser.parse_args()
#print(args)

#max_model_len = args.max_model_len
#gpu_memory_utilization = args.gpu_memory_utilization

# Init Model
#model_name = "Qwen/Qwen3-0.6B"
#model_name = "Qwen/Qwen3-1.7B"
model_name = "Qwen/Qwen3-4B"
#model_name = "/app/Qwen3-4B"
logger.info("Loading model: " + model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype="auto",
    device_map="auto"
)
logger.info("Loaded model: " + model_name)
tokenizer.model_max_length = 10000
logger.info("Model max context length: " + str(model.config.max_position_embeddings))


#====================================
#http://localhost:5008/run?task=Who+are+you%3F

app = Flask(__name__)

#run
@app.route("/run", methods=['GET'])
def run_get():
    task = request.args.get('task')
    if len(task) == 0:
        abort(500)
    else:
        logger.info('Task: ' + task)
        data = {
            "max_new_tokens": 500,
            "num_beams": 2,
#            "early_stopping": true,
            "max_time": 30,
            "task": task
        }
        results = run(data)
        return jsonify({'result': results})

#run
@app.route("/run", methods=['POST'])
def run_post():
    data = json.loads(request.data)
    results = run(data)
    return jsonify({'result': results})

#5008 port for LLM
if __name__ =='__main__':
#    app.run(host="0.0.0.0", port=5008, debug=False)
    app.run(host="0.0.0.0")
