from peft import PeftModel
import transformers
import gradio as gr

from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BloomForCausalLM, GenerationConfig
from transformers.models.opt.modeling_opt import OPTDecoderLayer

from flask import Flask, jsonify, request

app = Flask(__name__)

BASE_MODEL = 'bigscience/bloomz-7b1-mt'

LORA_WEIGHTS = "BLOOM_VI_test"




model = AutoModelForCausalLM.from_pretrained( 
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map={"":0},
)
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model.eval()
model = PeftModel.from_pretrained(model, LORA_WEIGHTS, torch_dtype=torch.float16)
model.half()
model = model.merge_and_unload() 
#model.save_pretrained("BLOOM_VI_test")
#model.push_to_hub("xieyang233/BLOOM_VI")

history = []

def evaluate(
    inputs,
    temperature=0.1,
    top_p=0.75,
    top_k=40,
    num_beams=4,
    max_new_tokens=256,
    **kwargs,    
):
    instruction = ""
    for pair in history:
        his_inputs = pair.get("inputs")
        his_outputs = pair.get("outputs")
        instruction += "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n#Instruction: {}\n#Response: {}\n".format(his_inputs, his_outputs)
    instruction += "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n#Instruction: {}\n#Response:".format(inputs)
    encodings = tokenizer(instruction, max_length=512, return_tensors="pt", truncation=True).to("cuda")
    generation_config = GenerationConfig(
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        num_beams=num_beams,
        do_sample=True,
        **kwargs,
    )    
    outputs = model.generate(input_ids=encodings["input_ids"], generation_config=generation_config, max_new_tokens=max_new_tokens)
    preds = tokenizer.decode(outputs[0])
    preds = preds.split("#Response:")[-1].strip().replace("</s>", "")
    item = {
        "inputs": inputs.strip(),
        "outputs": preds
    }
    history.append(item)

    return preds

    #prompt = generate_prompt(instruction, input)
    #inputs = tokenizer(prompt, return_tensors="pt")
    #input_ids = inputs["input_ids"].to(device)
    '''
    generation_config = GenerationConfig(
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        num_beams=num_beams,
        
        **kwargs,
    )

    with torch.no_grad():
        generation_output = model.generate(
            input_ids=input_ids,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True,
            max_new_tokens=max_new_tokens,
           
        )
    s = generation_output.sequences[0]
    output = tokenizer.decode(s)
    return output.split("#Response:")[1].strip()
    '''
@app.route('/api/evalute', methods=['POST'])
def api_evaluate():
    data = request.get_json()
    print('request data : \n', data)
    print()
    inputs = data['Input']
    response = evaluate(inputs)
    return jsonify(result=response)

# Old testing code follows.

'''
#"Tell me about the president of Mexico in 2019."
"Kể cho tôi về Tổng thống Mexico năm 2019.",
#"Tell me about the king of France in 2019."
"Kể cho tôi về vua của Pháp vào năm 2019.",
#"List all Canadian provinces in alphabetical order."
"Danh sách các tỉnh của Canada theo thứ tự bảng chữ cái.",
#"Write a Python program that prints the first 10 Fibonacci numbers."
"Viết một chương trình Python in ra 10 số Fibonacci đầu tiên.",
#"Write a program that prints the numbers from 1 to 100. But for multiples of three print 'Fizz' instead of the number and for the multiples of five print 'Buzz'. For numbers which are multiples of both three and five print 'FizzBuzz'."
"Viết một chương trình in ra các số từ 1 đến 100. Nhưng đối với các số chia hết cho ba, in ra 'Fizz' thay vì số đó và đối với các số chia hết cho năm, in ra 'Buzz'. Đối với các số chia hết cả cho ba và năm, in ra 'FizzBuzz'.",
#"Tell me five words that rhyme with 'shock'."
"Cho tôi biết năm từ có vần điệu với từ 'shock'.",
#"Translate the sentence 'I have no mouth but I must scream' into Spanish."
"Dịch câu 'I have no mouth but I must scream' sang tiếng Tây Ban Nha là",
#"Count up from 1 to 500."
"Đếm từ 1 đến 500."
'''
if __name__ == "__main__":
    # testing code for readme
    gr.Interface(
        fn=evaluate,
        inputs=[
            gr.components.Textbox(
                lines=2, label="Input", placeholder="Tell me about BLOOM-VI."
            ),
            gr.components.Slider(minimum=0, maximum=1, value=0.1, label="Temperature"),
            gr.components.Slider(minimum=0, maximum=1, value=0.75, label="Top p"),
            gr.components.Slider(minimum=0, maximum=100, step=1, value=40, label="Top k"),
            gr.components.Slider(minimum=1, maximum=4, step=1, value=4, label="Beams"),
            gr.components.Slider(
                minimum=1, maximum=2000, step=1, value=256, label="Max tokens"
            ),
        ],
        outputs=[
            gr.components.Textbox(
                lines=5,
                label="Output",
            )
        ],
        title="🌲 🌲 🌲 BLOOM-VI",
        description="BLOOM-VI is a 7b-parameter BLOOM model finetuned on Alpaca translated to Vietnamese.",
    ).launch(share=True)
    
    '''    
    for instruction in [
        #"Tell me about alpacas." 
        "Kể cho tôi về lạc đà Alpaca.",
        "Đếm từ 1 đến 500.",
        "Kể cho tôi về vua của Pháp vào năm 2019.",
        "Write a Python program that prints the first 10 Fibonacci numbers."
    ]:
        print("Instruction:", instruction)
        print("Response:", evaluate(instruction))
        print()
    '''