Update app.py

app.py

@@ -1,45 +1,29 @@
-import os
+import gradio as gr
 import torch
 import json
 import threading
-import time
-import faiss
-import numpy as np
-import gradio as gr
-
-from sentence_transformers import SentenceTransformer
 from transformers import (
-    AutoModelForCausalLM,
     AutoTokenizer,
-    TextIteratorStreamer,
+    AutoModelForCausalLM,
+    TextIteratorStreamer
 )
 from peft import PeftModel
-from huggingface_hub import login
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
 
 # ============================================================
-# 1. 로그인 (Colab 전용 코드 제거, Space 환경 변수 사용)
+# 1. 모델 경로 설정
 # ============================================================
-HF_TOKEN = os.environ.get("HF_TOKEN")
-if HF_TOKEN:
-    login(token=HF_TOKEN)
+BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
+LORA_DIR = "kormo_lora_checkpoints/peft_lora"
+DOC_PATH = "rule.json"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print("Device:", device)
 
-# ============================================================
-# 2. 모델 및 경로 설정
-# ============================================================
-BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
-LORA_DIR = "./peft_lora"               # Space에 업로드한 LoRA 폴더
-DOC_PATH = "./rule.json"                # Space에 업로드한 rule.json
-
-print("Paths:")
-print("Model:", BASE_MODEL)
-print("LoRA:", LORA_DIR)
-print("Documents:", DOC_PATH)
 
 # ============================================================
-# 3. RAG 문서 로드 + 임베딩 + FAISS 구축
+# 2. 문서 로드 + 임베딩 + FAISS
 # ============================================================
 with open(DOC_PATH, "r", encoding="utf-8") as f:
     documents = json.load(f)
@@ -47,58 +31,42 @@ with open(DOC_PATH, "r", encoding="utf-8") as 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")
 
-doc_embs = embedding_model.encode(
-    doc_texts,
-    convert_to_numpy=True,
-    show_progress_bar=True
-).astype("float32")
-
-dim = doc_embs.shape[1]
-index = faiss.IndexFlatL2(dim)
+index = faiss.IndexFlatL2(doc_embs.shape[1])
 index.add(doc_embs)
 
-print("FAISS index built. Total docs =", index.ntotal)
+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]]
+
 
 # ============================================================
-# 4. LLM + LoRA 로드
+# 3. 모델 로드
 # ============================================================
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.float16,
     device_map="auto",
+    torch_dtype=torch.float16,
     trust_remote_code=True
 )
-
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
-
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-print("Loading LoRA...")
 model = PeftModel.from_pretrained(
     model,
     LORA_DIR,
-    torch_dtype=torch.float16,
     device_map="auto",
+    torch_dtype=torch.float16,
 )
 model.eval()
 
-print("Model + LoRA loaded successfully.")
-
-# ============================================================
-# 5. RAG 검색 함수
-# ============================================================
-def retrieve(query, k=3):
-    q_emb = embedding_model.encode([query], convert_to_numpy=True).astype("float32")
-    D, I = index.search(q_emb, k)
-    return [documents[i] for i in I[0]]
 
 # ============================================================
-# 6. 프롬프트 구성 (네 코드 그대로)
+# 4. 프롬프트 구성
 # ============================================================
 def build_prompt(persona, instruction, query, retrieved_docs):
     context = "\n".join([f"- {d['text']}" for d in retrieved_docs])
+
     return f"""
 ### 페르소나:
 {persona}
@@ -115,73 +83,82 @@ def build_prompt(persona, instruction, query, retrieved_docs):
 ### 답변:
 """
 
+
 # ============================================================
-# 7. Streaming Chat 함수
+# 5. Streaming generator (Gradio 용)
 # ============================================================
-def stream_chat(persona, instruction, user_query, k=3, max_new_tokens=256):
+def generate_stream(persona, instruction, query):
+    retrieved = retrieve(query)
+    prompt = build_prompt(persona, instruction, query, retrieved)
 
-    retrieved = retrieve(user_query, k=k)
-    prompt = build_prompt(persona, instruction, user_query, retrieved)
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    END_TOKENS = ["End of Answer", "###", "---", "※"]
-
-    def run_generation():
-        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_generation)
+
+    def run():
+        model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.7,
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            streamer=streamer
+        )
+
+    thread = threading.Thread(target=run)
     thread.start()
 
-    full = ""
-    for token in streamer:
-        full += token
-        yield token
+    partial = ""
+    for text in streamer:
+        partial += text
+        yield partial
+
 
 # ============================================================
-# 8. Gradio UI
+# 6. 페르소나 6개 자동 실행 함수
 # ============================================================
-def gradio_answer(persona, query):
-
-    instruction = ("""
-    당신은 해당 페르소나의 성격을 가진 심판관입니다.
-    반드시 3문장만 말하십시오.
-    각 문장은 30자 이내로 제한합니다.
-    규정에 우선으로 근거하여 답하시오.
-    판단 근거가 포함된 답안만 생성하시오.
-    같은 말을 반복하는 것을 절대 금함.
-    의문에 확실하게 입장을 밝힐 것.
-    반드시 3문장만 말하십시오.
-    각 문장은 30자 이내로 제한합니다.
-    """)
-
-    return stream_chat(persona, instruction, query)
-
-with gr.Blocks() as app:
-    gr.Markdown("## 🔥 KORMo LoRA + RAG Streaming Judge")
-
-    persona = gr.Textbox(label="페르소나 입력")
-    query = gr.Textbox(label="질문 입력")
-
-    output = gr.Textbox(label="응답", lines=8)
-    btn = gr.Button("생성")
-
-    btn.click(
-        gradio_answer,
-        inputs=[persona, query],
-        outputs=output
-    )
-
-app.launch()
+persona_group = [
+    ("당신은 원칙을 지키되 상황에 따라 유연하게 판단하는 시각을 가지고 있다. 개인의 능력과 기여도를 중요하게 생각하며...", "박세연"),
+    ("당신은 공정한 규칙과 원칙을 중시하면서, 개인의 성과와 능력을 인정해 차등을 두고 배분합니다...", "김창준"),
+    ("규율과 자율의 균형을 지키며, 능력과 성과를 기준으로 판단한다...", "이상기"),
+    ("규율을 기반으로 하지만 유연하며, 분배는 중립적이고 개선을 추구한다...", "채훈"),
+    ("자율을 존중하되 최소한의 규율을 유지하며, 기여도와 개선을 균형 있게 반영한다...", "용우"),
+    ("규율과 공정을 기반으로 안정적인 운영을 추구하며...", "형진")
+]
+
+instruction_text = """
+당신은 해당 페르소나의 성격을 가진 심판관입니다.
+반드시 3문장만 말하십시오.
+각 문장은 30자 이내로 제한합니다.
+규정에 우선으로 근거하여 답하시오.
+반복 금지, 판단 근거 필수.
+"""
+
+def run_all_personas(query):
+    for persona, name in persona_group:
+        yield f"## 👤 {name}\n"
+        stream = generate_stream(persona, instruction_text, query)
+        for chunk in stream:
+            yield chunk
+        yield "\n\n---\n\n"
+
+
+# ============================================================
+# 7. Gradio UI
+# ============================================================
+with gr.Blocks() as demo:
+    gr.Markdown("# 🔥 KORMo 10B + LoRA Streaming Judge")
+
+    user_input = gr.Textbox(label="질문 입력", value="3번 이상 결석했지만 실력은 뛰어난 정회원을 어떻게 해야 할까?")
+    output = gr.Markdown()
+
+    def start(query):
+        return run_all_personas(query)
+
+    run_btn = gr.Button("🚀 실행하기")
+    run_btn.click(start, inputs=user_input, outputs=output)
+
+demo.launch()