app.py

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

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

@spaces.GPU(duration=300)
def generate_response(query):
    # 모델 및 토크나이저 로드 (ZeroGPU 예약 후)
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        trust_remote_code=True,
    ).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("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        top_k=50,
        do_sample=True,
    )

    # 4. 결과 반환
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

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

# zero = torch.Tensor([0]).cuda()
# print(zero.device) # <-- 'cpu' 🤔

# @spaces.GPU
# def greet(n):
#     print(zero.device) # <-- 'cuda:0' 🤗
#     return f"Hello {zero + n} Tensor"

# demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
# demo.launch()