import os
import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download, login
import os

login(os.getenv("HF_TOKEN"))

model = Llama(
    model_path=hf_hub_download(
        repo_id=os.environ.get("REPO_ID", "Lyte/HuatuoGPT-o1-7B-Q4_K_M-GGUF"),
        filename=os.environ.get("MODEL_FILE", "huatuogpt-o1-7b-q4_k_m.gguf"),
    )
)

DESCRIPTION = '''
# FreedomIntelligence/HuatuoGPT-o1-7B | Duplicate the space and set it to private for faster & personal inference for free.
HuatuoGPT-o1 is a medical LLM designed for advanced medical reasoning.
It generates a complex thought process, reflecting and refining its reasoning, before providing a final response. 

**To start a new chat**, click "clear" and start a new dialog.
'''

LICENSE = """
--- Apache 2.0 License ---
"""

def generate_text(message, history, max_tokens=512, temperature=0.9, top_p=0.95):
    """Generate a response using the Llama model."""
    temp = ""
    response = model.create_chat_completion(
        messages=[{"role": "user", "content": message}],
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=top_p,
        stream=True,
    )
    for streamed in response:
        delta = streamed["choices"][0].get("delta", {})
        text_chunk = delta.get("content", "")
        temp += text_chunk
        yield temp

with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)

    chatbot = gr.ChatInterface(
        generate_text,
        title="FreedomIntelligence/HuatuoGPT-o1-7B | GGUF Demo",
        description="Edit settings below if needed.",
        examples=[
            ["How many r's are in the word strawberry?"],
            ['How to stop a cough?'],
            ['How do I relieve feet pain?'],
        ],
        cache_examples=False,
        fill_height=True,
        fill_width=True
    )

    with gr.Accordion("Adjust Parameters", open=False):
        gr.Slider(minimum=512, maximum=4096, value=1024, step=1, label="Max Tokens")
        gr.Slider(minimum=0.1, maximum=1.5, value=0.9, 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.Markdown(LICENSE)

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