import os
import subprocess
from flask import Flask
from flask import request
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM 

name='Hello'
model_id = "Qwen/Qwen3-1.7B-Base"
# filename = "tinyllama-1.1b-chat-v1.0.Q6_K.gguf"


print(name, model_id)


# # pip install gguf
# import torch
# from transformers import AutoTokenizer, AutoModelForCausalLM


# dtype = torch.float32 # could be torch.float16 or torch.bfloat16 too
# tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)

# model = AutoModelForCausalLM.from_pretrained(model_id, gguf_file=filename, dtype=dtype)


# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
)

def get_texts(ask):
    
    # prepare the model input
    messages = [    {"role": "user", "content": ask},]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # conduct text completion
    generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=16384,
    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 

    content = tokenizer.decode(output_ids, skip_special_tokens=True)

    print("content:", content)
    return content



app = Flask(name)

@app.route('/')
def hello_world():
  return '<p>Hello, World!</p>'

@app.route('/gen/', methods=['GET'])
def gen_text():
    error = None
    query = request.args.get('ask')
    # проверяем, передается ли параметр
    # 'query' в URL-адресе
    if query and query != '':
        # если `query`существует и это не пустая строка,
        # то можно приступать к обработке запроса
        return f'<p>{get_texts(query)}</p>'
    else:
        # если `query` не существует или это пустая строка, то 
        # отображаем форму поискового запроса с сообщением.
        error = 'Не введен запрос!'
        return f'search.html'

if __name__ == '__main__':
  app.run(debug=True, host='0.0.0.0', port='7860')