"""
This is an example implementation of main.cpp from llama.cpp
Quirks:
 * Its not exactly alike since this port is designed around programmatic I/O
 * Input is always echoed if on, so it should be turned off when using "input()"
 * The first antiprompt should be the userprompt like "\nUser:", 
   because its added when n_predict is reached (aka generation ended prematurely)
 * n_predict can be set to -1 for unlimited length responses (or just a really high value)
 * Instruction mode adds its own antiprompt.
   You should also still be feeding the model with a "primer" prompt that 
   shows it the expected format.
"""

import ctypes
import sys
from time import time
from os import cpu_count, path

import llama_cpp
from common import GptParams, gpt_params_parse, gpt_random_prompt
import util


# A LLaMA interactive session
class LLaMAInteract:
    def __init__(self, params: GptParams) -> None:
        # input args
        self.params = params
        if self.params.path_session is None:
            self.params.path_session = ""
        if self.params.antiprompt is None:
            self.params.antiprompt = ""

        if self.params.perplexity:
            raise NotImplementedError(
                """************
please use the 'perplexity' tool for perplexity calculations
************"""
            )

        if self.params.embedding:
            raise NotImplementedError(
                """************
please use the 'embedding' tool for embedding calculations
************"""
            )

        if self.params.n_ctx > 2048:
            print(
                f"""warning: model does not support \
context sizes greater than 2048 tokens ({self.params.n_ctx} \
specified) expect poor results""",
                file=sys.stderr,
            )

        if self.params.seed <= 0:
            self.params.seed = int(time())

        print(f"seed = {self.params.seed}", file=sys.stderr)

        if self.params.random_prompt:
            self.params.prompt = gpt_random_prompt(self.params.seed)

        # runtime args
        self.input_consumed = 0
        self.n_past = 0
        self.n_session_consumed = 0
        self.first_antiprompt = []
        self.remaining_tokens = self.params.n_predict
        self.output_echo = self.params.input_echo
        self.multibyte_fix = []

        # model load
        self.lparams = llama_cpp.llama_model_default_params()
        self.lparams.n_ctx = self.params.n_ctx
        self.lparams.n_parts = self.params.n_parts
        self.lparams.seed = self.params.seed
        self.lparams.memory_f16 = self.params.memory_f16
        self.lparams.use_mlock = self.params.use_mlock
        self.lparams.use_mmap = self.params.use_mmap

        self.model = llama_cpp.llama_load_model_from_file(
            self.params.model.encode("utf8"), self.lparams
        )

        # Context Params.
        self.cparams = llama_cpp.llama_context_default_params()

        self.ctx = llama_cpp.llama_new_context_with_model(self.model, self.cparams)
        if not self.ctx:
            raise RuntimeError(f"error: failed to load model '{self.params.model}'")

        if self.params.ignore_eos:
            self.params.logit_bias[llama_cpp.llama_token_eos()] = -float("inf")

        if len(self.params.lora_adapter) > 0:
            if (
                llama_cpp.llama_apply_lora_from_file(
                    self.ctx,
                    self.params.lora_adapter.encode("utf8"),
                    (
                        self.params.lora_base.encode("utf8")
                        if len(self.params.lora_base) > 0
                        else None
                    ),
                    self.params.n_threads,
                )
                != 0
            ):
                print("error: failed to apply lora adapter")
                return

        print(file=sys.stderr)
        print(
            f"system_info: n_threads = {self.params.n_threads} / {cpu_count()} \
| {llama_cpp.llama_print_system_info().decode('utf8')}",
            file=sys.stderr,
        )

        # determine the required inference memory per token:
        if self.params.mem_test:
            tmp = [0, 1, 2, 3]
            llama_cpp.llama_eval(
                self.ctx,
                (llama_cpp.c_int * len(tmp))(*tmp),
                len(tmp),
                0,
                self.n_threads,
            )
            llama_cpp.llama_print_timings(self.ctx)
            self.exit()
            return

        # create internal context
        self.n_ctx = llama_cpp.llama_n_ctx(self.ctx)

        # Add a space in front of the first character to match OG llama tokenizer behavior
        self.params.prompt = " " + self.params.prompt

        # Load prompt file
        if self.params.file:
            with open(self.params.file) as f:
                self.params.prompt = f.read()

        self.session_tokens: list[llama_cpp.llama_token] = []
        if len(self.params.path_session) > 0:
            print(
                f"attempting to load saved session from '{self.params.path_session}'",
                file=sys.stderr,
            )

            if path.exists(self.params.path_session):
                _session_tokens = (llama_cpp.llama_token * (self.params.n_ctx))()
                _n_token_count_out = llama_cpp.c_size_t()
                if (
                    llama_cpp.llama_load_session_file(
                        self.ctx,
                        self.params.path_session.encode("utf8"),
                        _session_tokens,
                        self.params.n_ctx,
                        ctypes.byref(_n_token_count_out),
                    )
                    != 1
                ):
                    print(
                        f"error: failed to load session file '{self.params.path_session}'",
                        file=sys.stderr,
                    )
                    return
                _n_token_count_out = _n_token_count_out.value
                self.session_tokens = _session_tokens[:_n_token_count_out]
                print(
                    f"loaded a session with prompt size of {_n_token_count_out} tokens",
                    file=sys.stderr,
                )
            else:
                print(f"session file does not exist, will create", file=sys.stderr)

        # tokenize the prompt
        self.embd = []
        self.embd_inp = self._tokenize(self.params.prompt)

        if len(self.embd_inp) > self.n_ctx - 4:
            raise RuntimeError(
                f"error: prompt is too long ({len(self.embd_inp)} tokens, max {self.params.n_ctx - 4})"
            )

        # debug message about similarity of saved session, if applicable
        self.n_matching_session_tokens = 0
        if len(self.session_tokens) > 0:
            for id in self.session_tokens:
                if (
                    self.n_matching_session_tokens >= len(self.embd_inp)
                    or id != self.embd_inp[self.n_matching_session_tokens]
                ):
                    break
                self.n_matching_session_tokens += 1

            if self.n_matching_session_tokens >= len(self.embd_inp):
                print(f"session file has exact match for prompt!")
            elif self.n_matching_session_tokens < (len(self.embd_inp) / 2):
                print(
                    f"warning: session file has low similarity to prompt ({self.n_matching_session_tokens} / {len(self.embd_inp)} tokens); will mostly be reevaluated"
                )
            else:
                print(
                    f"session file matches {self.n_matching_session_tokens} / {len(self.embd_inp)} tokens of prompt"
                )

        self.need_to_save_session = len(
            self.params.path_session
        ) > 0 and self.n_matching_session_tokens < (len(self.embd_inp) * 3 / 4)

        # number of tokens to keep when resetting context
        if (
            self.params.n_keep < 0
            or self.params.n_keep > len(self.embd_inp)
            or self.params.instruct
        ):
            self.params.n_keep = len(self.embd_inp)

        self.inp_prefix = self._tokenize(self.params.instruct_inp_prefix)
        self.inp_suffix = self._tokenize(self.params.instruct_inp_suffix, False)

        # in instruct mode, we inject a prefix and a suffix to each input by the user
        self.antiecho = None
        if self.params.instruct:
            self.params.interactive_start = True
            _ptn = self._tokenize(self.params.instruct_inp_prefix.strip(), False)
            self.first_antiprompt.append(_ptn)
            self.antiecho = util.IterSearch(_ptn)

        # enable interactive mode if reverse prompt or interactive start is specified
        if len(self.params.antiprompt) != 0 or self.params.interactive_start:
            self.params.interactive = True

        # determine newline token
        self.llama_token_newline = self._tokenize("\n", False)
        self.llama_token_eot = self._tokenize(" [end of text]\n", False)

        if self.params.verbose_prompt:
            print(
                f"""
prompt: '{self.params.prompt}'
number of tokens in prompt = {len(self.embd_inp)}""",
                file=sys.stderr,
            )

            for i in range(len(self.embd_inp)):
                print(
                    f"{self.embd_inp[i]} -> '{self.token_to_str(self.embd_inp[i])}'",
                    file=sys.stderr,
                )

            if self.params.n_keep > 0:
                print("static prompt based on n_keep: '")
                for i in range(self.params.n_keep):
                    print(self.token_to_str(self.embd_inp[i]), file=sys.stderr)
                print("'", file=sys.stderr)
            print(file=sys.stderr)

        if self.params.interactive:
            print("interactive mode on.", file=sys.stderr)

            if len(self.params.antiprompt) > 0:
                for antiprompt in self.params.antiprompt:
                    print(f"Reverse prompt: '{antiprompt}'", file=sys.stderr)

            if len(self.params.input_prefix) > 0:
                print(f"Input prefix: '{self.params.input_prefix}'", file=sys.stderr)

        print(
            f"""sampling: repeat_last_n = {self.params.repeat_last_n},\
repeat_penalty = {self.params.repeat_penalty},\
presence_penalty = {self.params.presence_penalty},\
frequency_penalty = {self.params.frequency_penalty},\
top_k = {self.params.top_k},\
tfs_z = {self.params.tfs_z},\
top_p = {self.params.top_p},\
typical_p = {self.params.typical_p},\
temp = {self.params.temp},\
mirostat = {self.params.mirostat},\
mirostat_lr = {self.params.mirostat_eta},\
mirostat_ent = {self.params.mirostat_tau},\

generate: n_ctx = {self.n_ctx},\
n_batch = {self.params.n_batch},\
n_predict = {self.params.n_predict},\
n_keep = {self.params.n_keep}

""",
            file=sys.stderr,
        )

        # determine antiprompt tokens
        for i in self.params.antiprompt:
            self.first_antiprompt.append(self._tokenize(i, False))

        self.last_n_tokens = [0] * self.n_ctx  # TODO: deque doesnt support slices

        if params.interactive:
            print(
                """== Running in interactive mode. ==
 - Press Ctrl+C to interject at any time.
 - Press Return to return control to LLaMa.
 - If you want to submit another line, end your input in '\\'.

""",
                file=sys.stderr,
            )
        self.set_color(util.CONSOLE_COLOR_PROMPT)

    # tokenize a prompt
    def _tokenize(self, prompt, bos=True):
        _arr = (llama_cpp.llama_token * ((len(prompt) + 1) * 4))()
        _n = llama_cpp.llama_tokenize(
            self.model,
            prompt.encode("utf8", errors="ignore"),
            len(prompt),
            _arr,
            len(_arr),
            bos,
            False,
        )
        return _arr[:_n]

    def set_color(self, c):
        if self.params.use_color:
            print(c, end="")

    def use_antiprompt(self):
        return len(self.first_antiprompt) > 0

    # generate tokens
    def generate(self):
        while (
            self.remaining_tokens > 0
            or self.params.interactive
            or self.params.n_predict == -1
        ):
            # predict
            if len(self.embd) > 0:
                # infinite text generation via context swapping
                # if we run out of context:
                # - take the n_keep first tokens from the original prompt (via n_past)
                # - take half of the last (n_ctx - n_keep) tokens and recompute the logits in a batch
                if self.n_past + len(self.embd) > self.n_ctx:
                    n_left = self.n_past - self.params.n_keep
                    self.n_past = self.params.n_keep

                    # insert n_left/2 tokens at the start of embd from last_n_tokens
                    _insert = self.last_n_tokens[
                        self.n_ctx - int(n_left / 2) - len(self.embd) : -len(self.embd)
                    ]
                    self.embd = _insert + self.embd
                    self.params.path_session = ""

                # try to reuse a matching prefix from the loaded session instead of re-eval (via n_past)
                if self.n_session_consumed < len(self.session_tokens):
                    for i in range(len(self.embd)):
                        if self.embd[i] != self.session_tokens[self.n_session_consumed]:
                            self.session_tokens = self.session_tokens[
                                : self.n_session_consumed
                            ]
                            break

                        self.n_past += 1
                        self.n_session_consumed += 1

                        if self.n_session_consumed >= len(self.session_tokens):
                            i += 1
                            break

                    if i > 0:
                        self.embd = self.embd[i:]

                # evaluate tokens in batches
                # embd is typically prepared beforehand to fit within a batch, but not always
                # TODO BUG: The batching code causes nonsensical generation
                """for i in range(0, len(self.embd), self.params.n_batch):
					n_eval = self.params.n_batch
					_arr = (llama_cpp.llama_token * n_eval)(*self.embd[i:i + n_eval])
					if llama_cpp.llama_eval(self.ctx, _arr, n_eval, self.n_past, self.params.n_threads) != 0:
						print(f"failed to eval")
						return
					
					self.n_past += n_eval"""

                if (
                    llama_cpp.llama_eval(
                        self.ctx,
                        (llama_cpp.llama_token * len(self.embd))(*self.embd),
                        len(self.embd),
                        self.n_past,
                    )
                    != 0
                ):
                    raise Exception("Failed to llama_eval!")

                if len(self.embd) > 0 and len(self.params.path_session) > 0:
                    self.session_tokens.extend(self.embd)
                    self.n_session_consumed = len(self.session_tokens)

            self.n_past += len(self.embd)
            self.embd = []
            if len(self.embd_inp) <= self.input_consumed:  # && !is_interacting
                # out of user input, sample next token
                top_k = (
                    llama_cpp.llama_n_vocab(self.ctx)
                    if self.params.top_k <= 0
                    else self.params.top_k
                )
                repeat_last_n = (
                    self.n_ctx
                    if self.params.repeat_last_n < 0
                    else self.params.repeat_last_n
                )

                # optionally save the session on first sample (for faster prompt loading next time)
                if len(self.params.path_session) > 0 and self.need_to_save_session:
                    self.need_to_save_session = False
                    llama_cpp.llama_save_session_file(
                        self.ctx,
                        self.params.path_session.encode("utf8"),
                        (llama_cpp.llama_token * len(self.session_tokens))(
                            *self.session_tokens
                        ),
                        len(self.session_tokens),
                    )

                id = 0

                logits = llama_cpp.llama_get_logits(self.ctx)
                n_vocab = llama_cpp.llama_n_vocab(self.model)

                # Apply params.logit_bias map
                for key, value in self.params.logit_bias.items():
                    logits[key] += value

                _arr = (llama_cpp.llama_token_data * n_vocab)(
                    *[
                        llama_cpp.llama_token_data(token_id, logits[token_id], 0.0)
                        for token_id in range(n_vocab)
                    ]
                )
                candidates_p = llama_cpp.ctypes.pointer(
                    llama_cpp.llama_token_data_array(_arr, len(_arr), False)
                )

                # Apply penalties
                nl_logit = logits[llama_cpp.llama_token_nl(self.ctx)]
                last_n_repeat = min(len(self.last_n_tokens), repeat_last_n, self.n_ctx)

                _arr = (llama_cpp.llama_token * last_n_repeat)(
                    *self.last_n_tokens[len(self.last_n_tokens) - last_n_repeat :]
                )
                llama_cpp.llama_sample_repetition_penalties(
                    ctx=self.ctx,
                    candidates=candidates_p,
                    last_tokens_data=_arr,
                    penalty_last_n=last_n_repeat,
                    penalty_repeat=llama_cpp.c_float(self.params.repeat_penalty),
                    penalty_freq=llama_cpp.c_float(self.params.frequency_penalty),
                    penalty_present=llama_cpp.c_float(self.params.presence_penalty),
                )

                # NOT PRESENT IN CURRENT VERSION ?
                # llama_cpp.llama_sample_frequency_and_presence_penalti(self.ctx, candidates_p,
                # 	_arr,
                # 	last_n_repeat, llama_cpp.c_float(self.params.frequency_penalty), llama_cpp.c_float(self.params.presence_penalty))

                if not self.params.penalize_nl:
                    logits[llama_cpp.llama_token_nl()] = nl_logit

                if self.params.temp <= 0:
                    # Greedy sampling
                    id = llama_cpp.llama_sample_token_greedy(self.ctx, candidates_p)
                else:
                    if self.params.mirostat == 1:
                        mirostat_mu = 2.0 * self.params.mirostat_tau
                        mirostat_m = 100
                        llama_cpp.llama_sample_temperature(
                            self.ctx, candidates_p, llama_cpp.c_float(self.params.temp)
                        )
                        id = llama_cpp.llama_sample_token_mirostat(
                            self.ctx,
                            candidates_p,
                            llama_cpp.c_float(self.params.mirostat_tau),
                            llama_cpp.c_float(self.params.mirostat_eta),
                            llama_cpp.c_int(mirostat_m),
                            llama_cpp.c_float(mirostat_mu),
                        )
                    elif self.params.mirostat == 2:
                        mirostat_mu = 2.0 * self.params.mirostat_tau
                        llama_cpp.llama_sample_temperature(
                            self.ctx, candidates_p, llama_cpp.c_float(self.params.temp)
                        )
                        id = llama_cpp.llama_sample_token_mirostat_v2(
                            self.ctx,
                            candidates_p,
                            llama_cpp.c_float(self.params.mirostat_tau),
                            llama_cpp.c_float(self.params.mirostat_eta),
                            llama_cpp.c_float(mirostat_mu),
                        )
                    else:
                        # Temperature sampling
                        llama_cpp.llama_sample_top_k(
                            self.ctx,
                            candidates_p,
                            top_k,
                            min_keep=llama_cpp.c_size_t(1),
                        )
                        llama_cpp.llama_sample_tail_free(
                            self.ctx,
                            candidates_p,
                            llama_cpp.c_float(self.params.tfs_z),
                            min_keep=llama_cpp.c_size_t(1),
                        )
                        llama_cpp.llama_sample_typical(
                            self.ctx,
                            candidates_p,
                            llama_cpp.c_float(self.params.typical_p),
                            min_keep=llama_cpp.c_size_t(1),
                        )
                        llama_cpp.llama_sample_top_p(
                            self.ctx,
                            candidates_p,
                            llama_cpp.c_float(self.params.top_p),
                            min_keep=llama_cpp.c_size_t(1),
                        )
                        llama_cpp.llama_sample_temperature(
                            self.ctx, candidates_p, llama_cpp.c_float(self.params.temp)
                        )
                        id = llama_cpp.llama_sample_token(self.ctx, candidates_p)
                # print("`{}`".format(candidates_p.size))

                self.last_n_tokens.pop(0)
                self.last_n_tokens.append(id)

                # replace end of text token with newline token when in interactive mode
                if (
                    id == llama_cpp.llama_token_eos(self.ctx)
                    and self.params.interactive
                    and not self.params.instruct
                ):
                    id = self.llama_token_newline[0]
                    self.embd.append(id)
                    if self.use_antiprompt():
                        # tokenize and inject first reverse prompt
                        self.embd_inp += self.first_antiprompt[0]
                        for id in self.first_antiprompt[0]:
                            self.embd.append(id)
                else:
                    # add it to the context
                    self.embd.append(id)

                # echo this to console
                self.output_echo = True

                # decrement remaining sampling budget
                self.remaining_tokens -= 1
            else:
                # output to console if input echo is on
                self.output_echo = self.params.input_echo

                # some user input remains from prompt or interaction, forward it to processing
                while len(self.embd_inp) > self.input_consumed:
                    self.embd.append(self.embd_inp[self.input_consumed])
                    self.last_n_tokens.pop(0)
                    self.last_n_tokens.append(self.embd_inp[self.input_consumed])
                    self.input_consumed += 1
                    if len(self.embd) >= self.params.n_batch:
                        break

            # display tokens
            if self.output_echo:
                for id in self.embd:
                    if self.antiecho != None:
                        for r in self.antiecho(id):
                            yield r
                    else:
                        yield id

            # reset color to default if we there is no pending user input
            if self.params.input_echo and len(self.embd_inp) == self.input_consumed:
                self.set_color(util.CONSOLE_COLOR_DEFAULT)

            if self.params.interactive and len(self.embd_inp) <= self.input_consumed:
                # if antiprompt is present, stop
                if self.use_antiprompt():
                    if True in [
                        i == self.last_n_tokens[-len(i) :]
                        for i in self.first_antiprompt
                    ]:
                        break

                # if we are using instruction mode, and we have processed the initial prompt
                if self.params.interactive_start:
                    break

            # end of text token
            if len(self.embd) > 0 and self.embd[-1] == llama_cpp.llama_token_eos(
                self.ctx
            ):
                if not self.params.instruct:
                    for i in self.llama_token_eot:
                        yield i
                    break

            # respect n_predict even if antiprompt is present
            if (
                self.params.interactive
                and self.remaining_tokens <= 0
                and self.params.n_predict != -1
            ):
                # If we arent in instruction mode, fix the current generation by appending the antiprompt.
                # Makes it so if chat ends prematurely you dont append the AI's text etc.
                if not self.params.instruct:
                    self.embd_inp += self.first_antiprompt[0]
                self.n_remain = self.params.n_predict
                break

        self.params.interactive_start = False

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.exit()

    def exit(self):
        llama_cpp.llama_free(self.ctx)
        self.set_color(util.CONSOLE_COLOR_DEFAULT)

    def token_to_str(self, token_id: int) -> bytes:
        size = 32
        buffer = (ctypes.c_char * size)()
        n = llama_cpp.llama_token_to_piece(
            self.model, llama_cpp.llama_token(token_id), buffer, size
        )
        assert n <= size
        return bytes(buffer[:n])

    # return past text
    def past(self):
        for id in self.last_n_tokens[-self.n_past :]:
            yield self.token_to_str(id).decode("utf8", errors="ignore")

    # write input
    def input(self, prompt: str):
        if (
            self.params.instruct
            and self.last_n_tokens[-len(self.inp_prefix) :] != self.inp_prefix
        ):
            self.embd_inp += self.inp_prefix
        self.embd_inp += self._tokenize(prompt)
        if self.params.instruct:
            self.embd_inp += self.inp_suffix

    # write output
    def output(self):
        self.remaining_tokens = self.params.n_predict
        for id in self.generate():
            cur_char = self.token_to_str(id)

            # Add remainder of missing bytes
            if None in self.multibyte_fix:
                self.multibyte_fix[self.multibyte_fix.index(None)] = cur_char

            # Return completed utf char
            if len(self.multibyte_fix) > 0 and not None in self.multibyte_fix:
                yield (b"".join(self.multibyte_fix)).decode("utf8")
                self.multibyte_fix = []
                continue

            # Contains multi-byte UTF8
            for num, pattern in [(2, 192), (3, 224), (4, 240)]:
                # Bitwise AND check
                if pattern & int.from_bytes(cur_char, "little") == pattern:
                    self.multibyte_fix = [cur_char] + ([None] * (num - 1))

            # Stop incomplete bytes from passing
            if len(self.multibyte_fix) > 0:
                continue

            yield cur_char.decode("utf8")

    # read user input
    def read_input(self):
        out = ""
        while (t := input()).endswith("\\"):
            out += t[:-1] + "\n"
        return out + t + "\n"

    # interactive mode
    def interact(self):
        for i in self.output():
            print(i, end="", flush=True)
        self.params.input_echo = False

        # Using string instead of tokens to check for antiprompt,
        # It is more reliable than tokens for interactive mode.
        generated_str = ""
        while self.params.interactive:
            self.set_color(util.CONSOLE_COLOR_USER_INPUT)
            if self.params.instruct:
                print("\n> ", end="")
                self.input(self.read_input())
            else:
                print(self.params.input_prefix, end="")
                self.input(
                    f"{self.params.input_prefix}{self.read_input()}{self.params.input_suffix}"
                )
                print(self.params.input_suffix, end="")
            self.set_color(util.CONSOLE_COLOR_DEFAULT)

            try:
                for i in self.output():
                    print(i, end="", flush=True)
                    generated_str += i
                    for ap in self.params.antiprompt:
                        if generated_str.endswith(ap):
                            raise KeyboardInterrupt
            except KeyboardInterrupt:
                self.set_color(util.CONSOLE_COLOR_DEFAULT)
                if not self.params.instruct:
                    print(self.params.fix_prefix, end="")
                    self.input(self.params.fix_prefix)


if __name__ == "__main__":
    from datetime import datetime

    USER_NAME = "User"
    AI_NAME = "ChatLLaMa"

    time_now = datetime.now()
    prompt = f"""Text transcript of a never ending dialog, where {USER_NAME} interacts with an AI assistant named {AI_NAME}.
{AI_NAME} is helpful, kind, honest, friendly, good at writing and never fails to answer {USER_NAME}’s requests immediately and with details and precision.
Transcript below contains only the recorded dialog between two, without any annotations like (30 seconds passed...) or (to himself), just what {USER_NAME} and {AI_NAME} say aloud to each other.
The dialog lasts for years, the entirety of it is shared below. It's 10000 pages long.
The transcript only includes text, it does not include markup like HTML and Markdown.

{USER_NAME}: Hello, {AI_NAME}!
{AI_NAME}: Hello {USER_NAME}! How may I help you today?
{USER_NAME}: What time is it?
{AI_NAME}: It is {time_now.strftime("%H:%M")}.
{USER_NAME}: What year is it?
{AI_NAME}: We are in {time_now.strftime("%Y")}.
{USER_NAME}: What is a cat?
{AI_NAME}: A cat is a domestic species of small carnivorous mammal. It is the only domesticated species in the family Felidae.
{USER_NAME}: Name a color.
{AI_NAME}: Blue
{USER_NAME}:   """

    params = gpt_params_parse()
    if params.prompt is None and params.file is None:
        params.prompt = prompt

    with LLaMAInteract(params) as m:
        m.interact()
