import gradio as gr
import torch
import time
from transformers import LlamaForCausalLM, PreTrainedTokenizerFast, pipeline

model_name = "Azurro/APT3-1B-Instruct-v1"

tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name)
model = LlamaForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)

def generate_text(prompt, max_length, temperature, top_k, top_p):
    prompt = f'<s>[INST] {prompt.strip()} [/INST]'
    input_ids = tokenizer(prompt, return_tensors='pt', add_special_tokens=False).input_ids.to(model.device)
    start_time = time.time()
    output = model.generate(
        inputs=input_ids, 
        max_new_tokens=max_length,
        temperature=temperature, 
        top_k=top_k,
        do_sample=(temperature > 0),
        top_p=top_p,
        num_beams=1,
        bos_token_id=1,
        eos_token_id=2,
        pad_token_id=3,
        repetition_penalty=1.1
    )
    elapsed_time = time.time() - start_time
    decoded_output = tokenizer.decode(output[0])
    input_tokens_count = len(input_ids[0])
    input_chars_count = len(prompt)
    output_tokens_count = len(output[0])
    output_chars_count = len(decoded_output)
    gen_speed = output_tokens_count / elapsed_time
    decoded_output = decoded_output[len(prompt):].replace('</s>','').strip()
    print(f"Input tokens: {input_tokens_count} (chars: {input_chars_count}), Output tokens: {output_tokens_count} (chars: {output_chars_count}), Gen Time: {elapsed_time:.2f} secs ({gen_speed} toks/sec)")
    print(f"{'*'*10} Input {'*'*10}\n{prompt}")
    print(f"{'*'*10} Output {'*'*10}\n{prompt}")
    print(f"{'*'*30}")
    return decoded_output, input_tokens_count, input_chars_count, output_tokens_count, output_chars_count, gen_speed

demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.inputs.Textbox(label="Input Text"), 
        gr.inputs.Slider(1, 1000, step=1, default=100, label="Max Length"), 
        gr.inputs.Slider(0.0, 1.5, step=0.1, default=0.6, label="Temperature"), 
        gr.inputs.Slider(1, 400, step=1, default=200, label="Top K"), 
        gr.inputs.Slider(0.0, 1.0, step=0.05, default=0.95, label="Top P")
    ],
    outputs=[
        gr.outputs.Textbox(label="Generated Text"),
        gr.outputs.Textbox(label="Input Tokens Count"),
        gr.outputs.Textbox(label="Input Characters Count"),
        gr.outputs.Textbox(label="Output Tokens Count"),
        gr.outputs.Textbox(label="Output Characters Count"),
        gr.outputs.Textbox(label="Generation speed in tokens per second"),
    ]
)
demo.queue(concurrency_count=1)
demo.launch(max_threads=20)