import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from retriever.vectordb import search_documents  # 🧠 RAG 검색기 불러오기

model_name = "dasomaru/gemma-3-4bit-it-demo"


# 🚀 tokenizer는 CPU에서도 미리 불러올 수 있음
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# 🚀 model은 CPU로만 먼저 올림 (GPU 아직 없음)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # 4bit model이니까
    trust_remote_code=True,
)

@spaces.GPU(duration=300)
def generate_response(query):
    # 🚀 generate_response 함수 안에서 매번 로드
    # tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    # model = AutoModelForCausalLM.from_pretrained(
    #     model_name,
    #     torch_dtype=torch.float16,
    #     device_map="auto",  # ✅ 중요: 자동으로 GPU 할당
    #     trust_remote_code=True,
    # )
    tokenizer = AutoTokenizer.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
    model = AutoModelForCausalLM.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
    model.to("cuda")    

    # 1. 검색
    top_k = 5
    retrieved_docs = search_documents(query, top_k=top_k)

    # 2. 프롬프트 조립
    prompt = (
        "당신은 공인중개사 시험 문제 출제 전문가입니다.\n\n"
        "다음은 기출 문제 및 관련 법령 정보입니다:\n"
    )
    for idx, doc in enumerate(retrieved_docs, 1):
        prompt += f"- {doc}\n"
    prompt += f"\n이 정보를 참고하여 사용자의 요청에 답변해 주세요.\n\n"
    prompt += f"[질문]\n{query}\n\n[답변]\n"

    # 3. 답변 생성
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)  # ✅ model.device
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        top_k=50,
        do_sample=True,
    )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(fn=generate_response, inputs="text", outputs="text")
demo.launch()