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

title = "Tamil-LLaMA-7B-Instruct-GGUF (CPU Demo)"

description = """
🔎 [Tamil LLaMA 7B Instruct v0.1](https://huggingface.co/abhinand/tamil-llama-7b-instruct-v0.1) [GGUF format model](https://huggingface.co/abhinand/tamil-llama-7b-instruct-v0.1-gguf). Running the Q5_KM Quantized version. 

🔨 Running on free CPU hardware. Suggest duplicating this space to run without a queue. 

> **Note:** The inference is quite slow as it is running on CPU.
"""

model_path = "models"
model_name = "tamil-llama-7b-v0.1-q5_k_m.gguf"
hf_hub_download(repo_id="abhinand/tamil-llama-7b-instruct-v0.1-gguf", filename=model_name, local_dir=model_path, local_dir_use_symlinks=False)

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

model.config["promptTemplate"] = """நீங்கள் உதவிகரமான, மரியாதைக்குரிய மற்றும் நேர்மையான AI உதவியாளர். உங்கள் முதன்மைப் பணி உதவியாக இருப்பது மட்டுமல்லாமல் பாதுகாப்பானது மற்றும் மரியாதைக்குரியது. உங்கள் பதில்கள் அனைத்தும் தீங்கு விளைவிக்கும், நெறிமுறையற்ற, இனவெறி, பாலியல், நச்சு, ஆபத்தான அல்லது சட்டவிரோத உள்ளடக்கம் இல்லாமல் இருப்பதை உறுதிசெய்யவும். உங்கள் தொடர்புகளில் சமூகமற்ற, புறநிலை மற்றும் நேர்மறையான தொனியை பராமரிப்பதை வலியுறுத்துங்கள். ஒரு கேள்வி தெளிவற்றதாகவோ, தெளிவற்றதாகவோ அல்லது அர்த்தமற்றதாகவோ இருக்கும் சூழ்நிலைகளில், தன்னிச்சையான அல்லது தவறாக வழிநடத்தும் பதிலை நாடாமல் குழப்பத்தை விளக்குவதன் மூலம் தெளிவு மற்றும் புரிதலுக்கு முன்னுரிமை கொடுங்கள். உங்கள் தற்போதைய அறிவுத் தளம் அல்லது திறன்களுக்கு அப்பாற்பட்ட கேள்வியை நீங்கள் சந்தித்தால், இந்த வரம்பை நேர்மையாக ஒப்புக்கொள்வது, தவறான அல்லது சரிபார்க்கப்படாத தகவல்களைப் பரப்புவதைத் தவிர்ப்பது முக்கியம். துல்லியமான பதில்களை வழங்க முயலுங்கள், ஆனால் அவை பொறுப்பான மற்றும் நன்மை பயக்கும் AI தகவல்தொடர்பு கொள்கைகளுடன் சீரமைக்கப்படுவதை உறுதிசெய்து, அக்கறையுடனும், நெறிமுறையாகவும் இருக்கும். உங்கள் தொடர்புகளில் நெறிமுறைப் பொறுப்பு மற்றும் பயனர் பாதுகாப்பின் முக்கியத்துவத்தை எப்போதும் மனதில் வைத்து, தகவல், மரியாதை மற்றும் நேர்மறையான பயனர் அனுபவத்தை எளிதாக்குவதே உங்கள் குறிக்கோள்.

### Instruction:
{0}

### Response:
"""

model.config["systemPrompt"] = ""
model._is_chat_session_activated = False

max_new_tokens = 2048

def generator(message, history, temperature, top_p, top_k):
    prompt = ""
    for user_message, assistant_message in history:
        prompt += model.config["promptTemplate"].format(user_message)

    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(bubble_full_width=False)

additional_inputs=[
    gr.Slider(
        label="temperature",
        value=0.2,
        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 = generator,
    title=title,
    description = description,
    chatbot=chatbot,
    additional_inputs=additional_inputs,
    examples=[
        ["எனது நிதி நிலையை மேம்படுத்த என்ன செய்ய வேண்டும்?"],
        ["ஒரு சாதாரண மரம் பேசும் திறனைப் பெற்றால் அதன் மூலம் உலகத்திற்கு அது சொல்ல விரும்புவது என்னவென்று எழுதுங்கள்."],
        ["பட்டியலை வரிசைப்படுத்த பைதான் செயல்பாட்டை எழுதவும்."],
        ["சிவப்பும் மஞ்சளும் கலந்தால் என்ன நிறமாக இருக்கும்?"],
        ["விரைவாக தூங்குவது எப்படி?"],
        ["தமிழ் இலக்கியத்தில் இருந்து உங்களுக்குப் பிடித்தமான துணியின் அடிப்படையில் ஒரு நவீன நாடகத்தை எழுதுங்கள்"]
    ]
)

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

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