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="praveenpankaj/aksara_1_unsloth_q4",
        filename="aksara_-unsloth.Q4_K_M.gguf",
    ),
    n_ctx=1024,
) 

history = []

def generate_text(message, history):
    temp = ""
    input_prompt = "Ask akṣara anything on Agriculture in the Global South.\n"
    for interaction in history:
        input_prompt += "[|Umano|] " + interaction[0] + "\n"
        input_prompt += "[|Assistente|]" + interaction[1]
    
    input_prompt += "[|Umano|] " + message + "\n[|Assistente|]"

    print(input_prompt)

    output = llm(
        input_prompt,
        temperature=0.15,
        top_p=0.1,
        top_k=40, 
        repeat_penalty=1.1,
        max_tokens=1024,
        stop=[
            "[|Umano|]",
            "[|Assistente|]",
        ],
        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="akṣara running on CPU (quantized Q4_K)",
    description="This is a quantized version of akṣara running on CPU. It is a quantized version of the original version, that is running on a CPU machine.",
    examples=[
        "What are the recommended NPK dosage for maize varieties?",
        "Heavy rains are predicted next week. Is my rice crop ready for this, or should I harvest early?",
        "What crops can I grow during the dry season to use water more efficiently?"
    ],
    cache_examples=False,
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
)
demo.queue(concurrency_count=1, max_size=5)
demo.launch()