import gradio as gr
from llama_cpp import Llama

llm = Llama(
    model_path="mkllm-7b-instruct-q4_k_s.gguf",
    n_ctx=4096,            # Doubled context length
    n_threads=8,           # Increased threads
    n_batch=1024,          # Increased batch size
    use_mlock=True,
    use_mmap=True,
    n_gpu_layers=0,        # Set this to higher number if GPU available
    verbose=False          # Reduced logging for better performance
)

def format_chat_history(history):
    formatted_history = ""
    for user_msg, assistant_msg in history:
        formatted_history += f"USER: {user_msg}\nA: {assistant_msg}\n"
    return formatted_history

def chat(message, history):
    system_prompt = """Ти си МКЛЛМ-7Б, професионален асистент со вештачка интелигенција кој дава корисни информации.
ПРАВИЛА:
- Не измислувам податоци
- Давам точни информации
- Се фокусирам на фактите
- Одговарам ефикасно
- Одржувам професионален тон"""
    
    chat_history = format_chat_history(history)
    full_prompt = f"""SYSTEM: {system_prompt}
KONTEKST:
{chat_history}
USER: {message}
A: """
    
    response = llm(
        full_prompt,
        max_tokens=4096,    # Increased max tokens
        temperature=0.7,    # Keeping it precise
        top_p=0.1,
        repeat_penalty=1.2,
        top_k=20,
        stop=["USER:", "\n\n"],
        stream=True
    )
    
    partial_message = ""
    for chunk in response:
        if chunk and chunk['choices'][0]['text']:
            partial_message += chunk['choices'][0]['text']
            yield partial_message

demo = gr.ChatInterface(
    fn=chat,
    title="МКЛЛМ-7Б Стручен асистент",
    description="Професионален извор на информации, ВНИМАНИЕ, ЗНАЕ ДА ЛАЖЕ!!!",
    examples=[
        "Кои се основните принципи на базите на податоци?",
        "Објаснете како функционира HTTP протоколот?",
        "Кои се главните компоненти на електронскиот сметач и нивните функции?"
    ]
)

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )