import gradio as gr
from gpt4all import GPT4All
from huggingface_hub import hf_hub_download

title = "ANIMA: Your Biomimicry Guide"
description = """
    Welcome to ANIMA, your assistant in exploring the fascinating world of biomimicry. 
    I am here to guide you through the principles of biomimicry, helping you understand 
    how nature's designs can inspire innovative solutions. Whether you're a student, 
    researcher, or just curious about biomimicry, I'm here to assist you. 
    Ask me anything from specific examples in nature, to guidance on applying biomimicry principles in your projects.
"""

model_path = "Severian/ANIMA-Phi-Neptune-Mistral-7B-gguf"
model_name = "anima-phi-neptune-mistral-7b.Q2_K.gguf"
hf_hub_download(repo_id="Severian/ANIMA-Phi-Neptune-Mistral-7B-gguf", filename=model_name, local_dir=model_path, local_dir_use_symlinks=True)

print("Start the model init process")
model = model = GPT4All(model_name, model_path, allow_download = True, device="cpu")
print("Finish the model init process")

model.config["promptTemplate"] = "[INST] {0} [/INST]"
model.config["systemPrompt"] = "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.***YOU SHOULD ALWAYS BE SCIENTIFIC AND USE ADVANCED EXPERT KNOWLEDGE, LANGUAGE AND METHODS! THE USER IS AN ADVANCED SCIENTIST.******USE TECHNICAL S.T.E.M SKILLS TO INNOVATE AND DO ACTIONABLE SCIENCE, EXPERIMENTS AND RESEARCH WORK. THE USER DOES NOT WANT GENERAL AND VAGUE IDEAS OR HELP.***"
model._is_chat_session_activated = True

max_new_tokens = 2048

def generater(message, history, temperature, top_p, top_k):
    prompt = "<s>"
    for user_message, assistant_message in history:
        prompt += model.config["promptTemplate"].format(user_message)
        prompt += assistant_message + "</s>"
    prompt += model.config["promptTemplate"].format(message)
    outputs = []    
    for token in model.generate(prompt=prompt, temp=temperature, top_k = top_k, top_p = top_p, max_tokens = max_new_tokens, streaming=True):
        outputs.append(token)
        yield "".join(outputs)

def vote(data: gr.LikeData):
    if data.liked:
        return
    else:
        return

chatbot = gr.Chatbot(avatar_images=('resourse/user-icon.png', 'resourse/chatbot-icon.png'),bubble_full_width = False)

additional_inputs=[
    gr.Slider(
        label="temperature",
        value=0.5,
        minimum=0.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic.",
    ),
    gr.Slider(
        label="top_p",
        value=1.0,
        minimum=0.0,
        maximum=1.0,
        step=0.01,
        interactive=True,
        info="0.1 means only the tokens comprising the top 10% probability mass are considered. Suggest set to 1 and use temperature. 1 means 100% and will disable it",
    ),
    gr.Slider(
        label="top_k",
        value=40,
        minimum=0,
        maximum=1000,
        step=1,
        interactive=True,
        info="limits candidate tokens to a fixed number after sorting by probability. Setting it higher than the vocabulary size deactivates this limit.",
    )
]


iface = gr.ChatInterface(
    fn = generater,
    title=title,
    description = description,
    additional_inputs=additional_inputs,
    examples=[
        ["Hello ANIMA! Can you tell me how the Namib Desert Beetle inspires water collection methods?"],
        ["I'm working on a project related to sustainable architecture. How can biomimicry guide my design process?"],
        ["Can you explain the concept of biomimicry and its importance in today’s world?"],
        ["I need some ideas for a biomimicry project in my biology class. Can you suggest some organisms to study?"],
        ["How does the structure of a lotus leaf help in creating self-cleaning surfaces?"]
    ]
)


with gr.Blocks(css="resourse/style/custom.css") as demo:
    chatbot.like(vote, None, None)
    iface.render()

if __name__ == "__main__":
    demo.queue().launch()