import os
import gradio as gr

from text_generation import Client, InferenceAPIClient


def get_client(model: str):
    return InferenceAPIClient(model, token=os.getenv("HF_TOKEN", None),timeout=100)


def get_usernames(model: str):
    """
    Returns:
        (str, str, str, str): pre-prompt, username, bot name, separator
    """
    if model in ("OpenAssistant/oasst-sft-1-pythia-12b", "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"):
        return "", "<|prompter|>", "<|assistant|>", "<|endoftext|>"
    return "", "User: ", "Assistant: ", "\n"


def predict(
    inputs: str,
):
    model = "OpenAssistant/oasst-sft-1-pythia-12b"
    client = get_client(model)
    preprompt, user_name, assistant_name, sep = get_usernames(model)

    past = []
    limits = ",in max 200 words"
    total_inputs = preprompt + "".join(past) + inputs + limits + sep + assistant_name.rstrip() 

    partial_words = ""

    if model in ("OpenAssistant/oasst-sft-1-pythia-12b", "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"):
        iterator = client.generate(
            total_inputs,
            typical_p=0.1,
            truncate=1000,
            watermark=0,
            max_new_tokens=502,
        )


    yield iterator.generated_text

g = gr.Interface(
    fn=predict,
    inputs=[

        gr.components.Textbox(lines=3, label="Hi, how can I help you?", placeholder=""),
    ],
    outputs=[
        gr.inputs.Textbox(
            lines=10,
            label="",
        )
    ]
)
g.queue(concurrency_count=1)
g.launch()