import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 모델 로드
MODEL_NAME = "ewhk9887/deepseek-cokoder"

device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
model.to(device)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# 코드 리뷰 생성 함수
def generate_code_review(user_input, system_message, max_tokens, temperature, top_p):
    """
    사용자 입력과 시스템 메시지를 기반으로 AI 코드 리뷰 생성.
    """
    # 메시지 포맷 생성
    prompt = f"{system_message}\n\nCode:\n{user_input}\n\nReview:"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # 모델 출력 생성
    outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id
    )
    review = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return review

# Gradio 인터페이스 생성
with gr.Blocks() as demo:
    gr.Markdown("# DeepSeek Code Review Assistant")
    gr.Markdown("AI가 코드에 대한 리뷰를 제공합니다. 코드를 입력하고 리뷰를 확인하세요!")

    with gr.Row():
        with gr.Column():
            code_input = gr.Textbox(label="코드 입력", placeholder="리뷰를 원하는 코드를 입력하세요...", lines=10)
            system_message = gr.Textbox(
                label="시스템 메시지",
                value="유저의 코드에서 오류와 개선점을 한국어로 리뷰하세요.",
                lines=3,
            )
        with gr.Column():
            review_output = gr.Textbox(label="코드 리뷰 결과", lines=10)

    # 추가 옵션
    max_tokens = gr.Slider(label="Max Tokens", minimum=10, maximum=512, value=256, step=10)
    temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1)
    top_p = gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, value=0.9, step=0.05)

    # 버튼
    generate_button = gr.Button("리뷰 생성")

    # 이벤트 연결
    generate_button.click(
        fn=generate_code_review,
        inputs=[code_input, system_message, max_tokens, temperature, top_p],
        outputs=review_output,
    )

# 실행
if __name__ == "__main__":
    demo.launch()