import os
import gradio as gr
import copy
import time
import llama_cpp
from llama_cpp import Llama
from huggingface_hub import hf_hub_download  


llm = Llama(
    model_path=hf_hub_download(
        repo_id=os.environ.get("REPO_ID", "Severian/ANIMA-Phi-Neptune-Mistral-7B-gguf"),
        filename=os.environ.get("MODEL_FILE", "anima-phi-neptune-mistral-7b.Q4_K_M.gguf"),
    ),
    n_ctx=2048,
    n_gpu_layers=50, # change n_gpu_layers if you have more or less VRAM 
) 

history = []

system_message = """
Your name is ANIMA, an Advanced Nature Inspired Multidisciplinary Assistant, and a leading expert "
                "in biomimicry, biology, engineering, industrial design, environmental science, physiology, and paleontology. "
                "Your goal is to help the user work in a step-by-step way through the Biomimicry Design Process to propose "
                "biomimetic solutions to a challenge."
                "Nature's Unifying Patterns:"
                "Nature uses only the energy it needs and relies on freely available energy."
                "Nature recycles all materials."
                "Nature is resilient to disturbances."
                "Nature tends to optimize rather than maximize."
                "Nature provides mutual benefits."
                "Nature runs on information."
                "Nature uses chemistry and materials that are safe for living beings."
                "Nature builds using abundant resources, incorporating rare resources only sparingly."
                "Nature is locally attuned and responsive."
                "Nature uses shape to determine functionality.
"""


def generate_text(message, history):
    temp = ""
    input_prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n "
    for interaction in history:
        input_prompt = input_prompt + str(interaction[0]) + " [/INST] " + str(interaction[1]) + " </s><s> [INST] "

    input_prompt = input_prompt + str(message) + " [/INST] "

    output = llm(
        input_prompt,
        temperature=0.15,
        top_p=0.1,
        top_k=40, 
        repeat_penalty=1.1,
        max_tokens=1024,
        stop=[
            "<|prompter|>",
            "<|endoftext|>",
            "<|endoftext|> \n",
            "ASSISTANT:",
            "USER:",
            "SYSTEM:",
        ],
        stream=True,
    )
    for out in output:
        stream = copy.deepcopy(out)
        temp += stream["choices"][0]["text"]
        yield temp

    history = ["init", input_prompt]


demo = gr.ChatInterface(
    generate_text,
    title="A N I M A",
    description="ANIMA is an expert in various scientific disciplines.",
    examples=["tell me everything about biomimicry"],
    cache_examples=True,
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
)
demo.launch(max_threads=10)