import gradio as gr
from ctransformers import AutoTokenizer, AutoModelForCausalLM
import sys 

# model id 
model_id = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"
model_file = "mistral-7b-instruct-v0.1.Q4_K_M.gguf"
model_type = "mistral"

# it's a quantization model
quant_model = AutoModelForCausalLM.from_pretrained(model_id, model_file = model_file , model_type= model_type)


def lechat_respond(
    message,
    history: list[tuple[str, str]],
    max_tokens,
    temperature,
    top_p,
    top_k
):
    # so mistral's instruct format . here i didn't use chat history cuase of computation 🙄
    text = f"""<s>[INST] {message} [/INST]"""   

    response = ""
    
    for next_token in quant_model(text, 
                                  max_new_tokens = int(max_tokens), 
                                  temperature = temperature, 
                                  top_p = top_p, 
                                  top_k = top_k,
                                  stream = True):
        response += next_token
        yield response


#chat interface for le_chat
mistral_chat = gr.ChatInterface(
    fn = lechat_respond,
    type = 'messages',
    # chatbot = gr.Chatbot(placeholder = "<h5>LLM running on cpu so it may take long time to respond to ur prompt !</h5>"),
    # textbox = gr.Textbox(placeholder= "Ask whatever", scale = 7, container = False), 
    additional_inputs=[
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.8, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1,maximum=1.0,value=0.95,step=0.05,label="Top-p (nucleus sampling)"),
        gr.Slider(minimum = 40, maximum = 10000, value = 40, step = 10, label = "Top-k"),
    ],
    # theme= "ocean",
    # examples= [["Write a haiku about destruction of human's and the raise of AI"], ["Which species will rule the Earth in the future"]],
    # cache_examples = True
)

if __name__ == "__main__":
    # print version for debugging in your HF Spaces logs
    print("Gradio v", gr.__version__, file=sys.stderr)

    mistral_chat.launch(
        server_name="0.0.0.0",     # listen on all interfaces
        mcp_server=True,           # <-- expose /gradio_api/mcp/sse for MCP clients 🔥
    )