app.py

# app.py
import json
import threading
import gradio as gr
import torch
import faiss
import numpy as np

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from peft import PeftModel

# -----------------------------
# 0. 환경 설정
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Device:", device)

# -----------------------------
# 1. 모델 / 경로 설정
# -----------------------------
BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
LORA_DIR = "peft_lora"
DOC_PATH = "rule.json"

# -----------------------------
# 2. RAG 문서 로드 + FAISS 준비
# -----------------------------
with open(DOC_PATH, "r", encoding="utf-8") as f:
    documents = json.load(f)

doc_texts = [d["text"] for d in documents]

embedding_model = SentenceTransformer("jhgan/ko-sroberta-multitask", device=device)
doc_embs = embedding_model.encode(doc_texts, convert_to_numpy=True).astype("float32")

index = faiss.IndexFlatL2(doc_embs.shape[1])
index.add(doc_embs)

def retrieve(query, k=3):
    q = embedding_model.encode([query], convert_to_numpy=True).astype("float32")
    D, I = index.search(q, k)
    return [documents[i] for i in I[0]]

print("FAISS ready, docs:", index.ntotal)

# -----------------------------
# 3. 토크나이저 + 모델 로드 (LoRA 포함)
# -----------------------------
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

model = PeftModel.from_pretrained(
    model,
    LORA_DIR,
    device_map="auto",
    torch_dtype=torch.float16
)
model.eval()
print("Model + LoRA loaded")

# -----------------------------
# 4. 프롬프트 빌더
# -----------------------------
def build_prompt(persona, instruction, query, retrieved_docs):
    context = "\n".join([f"- {d['text']}" for d in retrieved_docs])
    return f"""
### 페르소나:
{persona}

### 참고사항:
{instruction}

### 규정:
{context}

### 질문:
{query}

### 답변:
"""

# -----------------------------
# 5. 스트리밍 UI용
# -----------------------------
def generate_stream(persona, instruction, query, max_new_tokens=256):
    retrieved = retrieve(query, k=3)
    prompt = build_prompt(persona, instruction, query, retrieved)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    def run_generate():
        with torch.no_grad():
            model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                top_p=0.9,
                temperature=0.7,
                repetition_penalty=1.2,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                streamer=streamer,
                use_cache=True
            )

    thread = threading.Thread(target=run_generate)
    thread.start()

    accumulated = ""
    for token in streamer:
        accumulated += token
        yield accumulated

# -----------------------------
# 6. API용 동기 생성
# -----------------------------
def generate_once(persona, instruction, query, max_new_tokens=256):
    retrieved = retrieve(query, k=3)
    prompt = build_prompt(persona, instruction, query, retrieved)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            top_p=0.9,
            temperature=0.7,
            repetition_penalty=1.2,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            use_cache=True
        )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return text.replace(prompt, "").strip()

# -----------------------------
# 7. 페르소나 그룹
# -----------------------------
persona_group = [
    ("당신은 원칙을 지키되 상황에 따라 유연하게 판단하는 시각을 가지고 있다. 개인의 능력과 기여도를 중요하게 생각하여 성과에 따른 차등 대우를 정당하다고 판단하며, 변화와 혁신을 최우선으로 여겨 관습보다 개선을 선택한다. 또한 내부에 머무르기보다 외부와의 연계와 협업을 적극적으로 추구하며, 학술 활동과 친목 활동의 균형을 통해 건강한 공동체 문화를 지향한다. 대외적으로 보여줄 수 있는 확실한 성과와 완성도를 중시하면서도, 단기적 해결과 장기적 기반 마련 사이에서 균형을 유지하려 노력한다.", '박세연'),
    ("당신은 공정한 규칙과 원칙을 중시하면서, 개인의 성과와 능력을 인정해 차등을 두고 배분합니다. 전통을 존중하되 점진적인 개선을 수용하며, 내부 활동에 머무르지 않고 외부와의 협업과 네트워크를 적극적으로 추구합니다. 회원 간 유대와 즐거움을 중요시하고, 완성도 높은 결과물과 과정에서의 배움을 모두 중시하며, 당장의 문제 해결과 장기적 기반 구축을 동시에 고려합니다.",'김창준'),
    ("규율과 자율의 균형을 지키며, 능력과 성과를 기준으로 판단한다. 전통을 유지하되 점진적 개선을 추구하고, 외부와의 협업을 적극적으로 모색한다. 즐거운 분위기 속에서 학습하며 개인의 성장을 중시하고, 단기 성과보다 동아리의 장기적 기반을 우선한다.", '이상기'),
    ("규율을 기반으로 하지만 유연하며, 분배는 중립적이고 개선을 추구한다. 외부 연계를 적당히 활용하며 학술·친목 모두 상황에 따라 선택하고, 가시성과 장기 기반을 조화롭게 고려한다.", '채훈'),
    ("자율을 존중하되 최소한의 규율을 유지하며, 기여도와 개선을 균형 있게 반영한다. 내부와 외부 활동을 상황에 따라 조절하고 학술과 친목 모두를 포용하며, 성장과 장기 기반을 중시하는 실용적 운영을 선호한다.", '용우'),
    ("규율과 공정을 기반으로 안정적인 운영을 추구하며, 균등·개선·친목 간의 균형을 중시한다. 내부 중심이되 필요에 따라 외부 협력을 수용하고, 성장과 장기 기반을 함께 고려하는 실용적 판단을 지향한다.",'형진')
]

instruction_text = """
해당 페르소나의 성격을 가진 심판관입니다.
반드시 3문장만 말하십시오.
각 문장은 30자 이내로 제한합니다.
규정에 근거하여 답하시오.
반복 금지, 판단 근거 필수.
"""

# -----------------------------
# 8. 스트리밍 UI용
# -----------------------------
def run_all_streaming(query):
    for persona, name in persona_group:
        yield f"## 👤 {name}\n"
        for partial in generate_stream(persona, instruction_text, query):
            yield partial
        yield "\n\n---\n\n"

# -----------------------------
# 9. API용 동기 실행 (문자열 반환)
# -----------------------------
def run_all_api(query):
    out = ""
    for persona, name in persona_group:
        out += f"## 👤 {name}\n"
        text = generate_once(persona, instruction_text, query)
        out += text + "\n\n---\n\n"
    return out

# -----------------------------
# 10. Gradio 앱 구성
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🔥 KORMo LoRA + RAG (Streaming UI + API)")

    user_input = gr.Textbox(
        label="질문 입력",
        value="3번 이상의 결석을 했지만 실력은 뛰어난 정회원을 어떻게 해야 할까?"
    )

    # Streaming UI
    output_stream = gr.Markdown()
    run_btn = gr.Button("🚀 실행(Streaming UI)")
    run_btn.click(fn=run_all_streaming, inputs=[user_input], outputs=[output_stream])

    # API 버튼 (동기 반환)
    api_output = gr.Textbox(label="API 반환 결과", lines=15)
    run_btn_api = gr.Button("🔁 실행(API)")
    run_btn_api.click(fn=run_all_api, inputs=[user_input], outputs=[api_output], api_name="start_api")

# -----------------------------
# 11. Launch
# -----------------------------
demo.launch(share=True)