import spaces

import os
import json
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer


@spaces.GPU()
def predict(message, history, system_prompt, temperature, max_tokens):
    messages = [{"role": "system", "content": system_prompt}]
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    stop_tokens = ["<|im_end|>", "<|endoftext|>", "<|im_start|>"]
    sampling_params = SamplingParams(temperature=temperature, top_p=1, max_tokens=max_tokens, stop=stop_tokens)
    completions = llm.generate(prompt, sampling_params)
    for output in completions:
        prompt = output.prompt
        print('==========================question=============================')
        print(prompt)
        generated_text = output.outputs[0].text
        print('===========================answer=============================')
        print(generated_text)
        for idx in range(len(generated_text)):
                yield generated_text[:idx+1]


if __name__ == "__main__":
    path = "stabilityai/stablelm-2-12b-chat"
    tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
    llm = LLM(model=path, tensor_parallel_size=1, trust_remote_code=True)
    gr.ChatInterface(
        predict,
        title="LLM playground",
        description="This is a LLM playground for StableLM",
        theme="soft",
        chatbot=gr.Chatbot(height=1400, label="Chat History",),
        textbox=gr.Textbox(placeholder="input", container=False, scale=7),
        retry_btn=None,
        undo_btn="Delete Previous",
        clear_btn="Clear",
        additional_inputs=[
            gr.Textbox("You are a hepful assistant.", label="System Prompt"),
            gr.Slider(0, 1, 0.7, label="Temperature"),
            gr.Slider(100, 2048, 1024, label="Max Tokens"),
        ],
        additional_inputs_accordion_name="Parameters",
        examples=[
            ["implement snake game using pygame"],
            ["Can you explain briefly to me what is the Python programming language?"],
            ["write a program to find the factorial of a number"],
        ],
    ).queue().launch()