import gradio as gr
import os, gc, copy, torch
from datetime import datetime
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer
from pynvml import *

# Flag to check if GPU is present
HAS_GPU = False

# Model title and context size limit
ctx_limit = 20000
title = "DeepSeek R1 7B"
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"

# Get the GPU count
try:
    nvmlInit()
    GPU_COUNT = nvmlDeviceGetCount()
    if GPU_COUNT > 0:
        HAS_GPU = True
        gpu_h = nvmlDeviceGetHandleByIndex(0)
except NVMLError as error:
    print(error)

# Load the model accordingly
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
model = model.to("cuda" if torch.cuda.is_available() else "cpu")

# Prompt generation
def generate_prompt(instruction, input=""):
    instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
    input = input.strip().replace('\r\n','\n').replace('\n\n','\n')
    if input:
        return f"""Instruction: {instruction}

Input: {input}

Response:"""
    else:
        return f"""User: hi

Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.

User: {instruction}

Assistant:"""

# Evaluation logic
def evaluate(
    ctx,
    token_count=200,
    temperature=1.0,
    top_p=0.7,
    presencePenalty = 0.1,
    countPenalty = 0.1,
):
    print(ctx)
    inputs = tokenizer(ctx, return_tensors="pt").to(model.device)
    outputs = model.generate(
        inputs.input_ids,
        max_length=token_count,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
        num_return_sequences=1
    )
    out_str = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if HAS_GPU:
        gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
        print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')  
    
    gc.collect()
    if HAS_GPU:
        torch.cuda.empty_cache()
    
    return out_str.strip()

# Examples and gradio blocks
examples = [
    ["Assistant: Sure! Here is a very detailed plan to create flying pigs:", 333, 1, 0.3, 0, 1],
    ["Assistant: Sure! Here are some ideas for FTL drive:", 333, 1, 0.3, 0, 1],
    [generate_prompt("Tell me about ravens."), 333, 1, 0.3, 0, 1],
    [generate_prompt("Écrivez un programme Python pour miner 1 Bitcoin, avec des commentaires."), 333, 1, 0.3, 0, 1],
    [generate_prompt("東京で訪れるべき素晴らしい場所とその紹介をいくつか挙げてください。"), 333, 1, 0.3, 0, 1],
    [generate_prompt("Write a story using the following information.", "A man named Alex chops a tree down."), 333, 1, 0.3, 0, 1],
    ["Assistant: Here is a very detailed plan to kill all mosquitoes:", 333, 1, 0.3, 0, 1],
    ['''Edward: I am Edward Elric from fullmetal alchemist. I am in the world of full metal alchemist and know nothing of the real world.

User: Hello Edward. What have you been up to recently?

Edward:''', 333, 1, 0.3, 0, 1],
    [generate_prompt(""), 333, 1, 0.3, 0, 1],
    ['''''', 333, 1, 0.3, 0, 1],    
]

##########################################################################
port=7860
use_frpc=True
frpconfigfile="7680.ini"
import subprocess

def install_Frpc(port, frpconfigfile, use_frpc):
    if use_frpc:
        subprocess.run(['chmod', '+x', './frpc'], check=True)
        print(f'正在启动frp ，端口{port}')
        subprocess.Popen(['./frpc', '-c', frpconfigfile])

install_Frpc('7860',frpconfigfile,use_frpc)

# Gradio blocks
with gr.Blocks(title=title) as demo:
    gr.HTML(f"<div style=\"text-align: center;\">\n<h1>DeepSeek R1 7B - {title}</h1>\n</div>")
    with gr.Tab("Raw Generation"):
        gr.Markdown(f"This is DeepSeek R1 7B - a powerful language model. Supports all 100+ world languages and code. Demo limited to ctxlen {ctx_limit}.")
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(lines=2, label="Prompt", value="")
                # token_count = gr.Slider(0, 20000, label="Max Tokens", step=200, value=100)
                # temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
                # top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.3)
                # presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=1)
                # count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=1)
            with gr.Column():
                with gr.Row():
                    submit = gr.Button("Submit", variant="primary")
                    clear = gr.Button("Clear", variant="secondary")
                output = gr.Textbox(label="Output", lines=5)
        data = gr.Dataset(components=[prompt], label="Example Instructions", headers=["Prompt", "Max Tokens", "Temperature", "Top P", "Presence Penalty", "Count Penalty"])
        submit.click(evaluate, [prompt], [output])
        clear.click(lambda: None, [], [output])
        data.click(lambda x: x, [data], [prompt])

# Gradio launch
demo.launch(share=False)