Update app.py

app.py

@@ -1,35 +1,42 @@
-import gradio as gr
-import torch
+# app.py
+import os
 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
+    TextIteratorStreamer,
 )
 from peft import PeftModel
-from sentence_transformers import SentenceTransformer
-import faiss
-import numpy as np
-
-# ============================================================
-# 1. 모델 경로 설정
-# ============================================================
-BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
-LORA_DIR = "peft_lora"
-DOC_PATH = "rule.json"
 
+# -----------------------------
+# 0. 환경 검사
+# -----------------------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-# ============================================================
-# 2. 문서 로드 + 임베딩 + FAISS
-# ============================================================
+print("Device:", device)
+
+# -----------------------------
+# 1. 모델 / 경로 설정
+# -----------------------------
+BASE_MODEL = "KORMo-Team/KORMo-10B-sft"   # 예시
+LORA_DIR = "peft_lora"                    # Space 내 업로드된 LoRA 폴더(또는 로컬 경로)
+DOC_PATH = "rule.json"                    # Space 내 업로드된 규정 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")
 
@@ -41,17 +48,24 @@ def retrieve(query, k=3):
     D, I = index.search(q, k)
     return [documents[i] for i in I[0]]
 
+print("FAISS ready, docs:", index.ntotal)
 
-# ============================================================
-# 3. 모델 로드
-# ============================================================
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+# -----------------------------
+# 3. 토크나이저·모델 로드 (LoRA 포함)
+# -----------------------------
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# 본체 모델 (device_map="auto" 사용하면 accelerate가 자동 분배)
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     device_map="auto",
     torch_dtype=torch.float16,
-    trust_remote_code=True
+    trust_remote_code=True,
 )
+
+# LoRA (PEFT) 적용
 model = PeftModel.from_pretrained(
     model,
     LORA_DIR,
@@ -59,14 +73,13 @@ model = PeftModel.from_pretrained(
     torch_dtype=torch.float16,
 )
 model.eval()
+print("Model + LoRA loaded")
 
-
-# ============================================================
-# 4. 프롬프트 구성
-# ============================================================
+# -----------------------------
+# 4. 프롬프트 빌더 (원본 그대로)
+# -----------------------------
 def build_prompt(persona, instruction, query, retrieved_docs):
     context = "\n".join([f"- {d['text']}" for d in retrieved_docs])
-
     return f"""
 ### 페르소나:
 {persona}
@@ -83,43 +96,68 @@ def build_prompt(persona, instruction, query, retrieved_docs):
 ### 답변:
 """
 
-
-# ============================================================
-# 5. Streaming generator (Gradio 용)
-# ============================================================
-def generate_stream(persona, instruction, query):
-    retrieved = retrieve(query)
+# -----------------------------
+# 5. 스트리밍 generator (UI용)
+#    - TextIteratorStreamer + 스레드 방식
+# -----------------------------
+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():
-        model.generate(
+    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  # Gradio의 스트리밍 출력은 부분 문자열을 계속 받게 함
+
+# -----------------------------
+# 6. 동기 생성 (API용) — 전체 텍스�� 반환
+#    - model.generate를 블록킹으로 실행하고 결과를 디코드
+# -----------------------------
+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=256,
+            max_new_tokens=max_new_tokens,
             do_sample=True,
             top_p=0.9,
             temperature=0.7,
             repetition_penalty=1.2,
-            pad_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
-            streamer=streamer
+            use_cache=True
         )
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # prompt 포함된 경우 제거
+    return text.replace(prompt, "").strip()
 
-    thread = threading.Thread(target=run)
-    thread.start()
-
-    partial = ""
-    for text in streamer:
-        partial += text
-        yield partial
-
-
-# ============================================================
-# 6. 페르소나 6개 자동 실행 함수
-# ============================================================
+# -----------------------------
+# 7. 페르소나 그룹 (원본 유지)
+# -----------------------------
 persona_group = [
     ("당신은 원칙을 지키되 상황에 따라 유연하게 판단하는 시각을 가지고 있다. 개인의 능력과 기여도를 중요하게 생각하며...", "박세연"),
     ("당신은 공정한 규칙과 원칙을 중시하면서, 개인의 성과와 능력을 인정해 차등을 두고 배분합니다...", "김창준"),
@@ -137,29 +175,51 @@ instruction_text = """
 반복 금지, 판단 근거 필수.
 """
 
-def run_all_personas(query):
+# -----------------------------
+# 8. UI용: 모든 페르소나에 대해 스트리밍 출력 (Gradio Blocks)
+# -----------------------------
+def run_all_streaming(query):
+    # Gradio에 문자열을 부분적으로 보여주고 싶을 때 yield를 사용
     for persona, name in persona_group:
-        yield f"## 👤 {name}\n"
-        stream = generate_stream(persona, instruction_text, query)
-        for chunk in stream:
-            yield chunk
+        header = f"## 👤 {name}\n"
+        yield header  # persona header
+        # streaming generator yields partials; 그걸 그대로 UI로 보냄
+        for partial in generate_stream(persona, instruction_text, query):
+            yield partial
         yield "\n\n---\n\n"
 
-
-# ============================================================
-# 7. Gradio UI
-# ============================================================
+# -----------------------------
+# 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 10B + LoRA Streaming Judge")
+    gr.Markdown("# 🔥 KORMo LoRA + RAG (Streaming UI + API)")
+    user_input = gr.Textbox(label="질문 입력", value="3번 이상의 결석을 했지만 실력은 동아리에서 뛰어난 정회원을 어떻게 해야 할까?")
+    output_stream = gr.Markdown()  # streaming UI에선 Markdown으로 실시간 갱신이 깔끔함
+
+    run_btn = gr.Button("🚀 실행(Streaming UI)")
+    run_btn.click(fn=run_all_streaming, inputs=[user_input], outputs=[output_stream])
 
-    user_input = gr.Textbox(label="질문 입력", value="3번 이상 결석했지만 실력은 뛰어난 정회원을 어떻게 해야 할까?")
-    output = gr.Markdown()
+    # API용 버튼 (비주얼용; 실제 API는 아래에 api_name으로 등록)
+    run_btn_api = gr.Button("🔁 실행(API, 동기)")
+    api_output = gr.Textbox(label="API 반환 결과", lines=10)
+    run_btn_api.click(fn=run_all_api, inputs=[user_input], outputs=[api_output])
 
-    def start(query):
-        for chunk in run_all_personas(query):
-            yield chunk
-    
-    run_btn = gr.Button("🚀 실행하기")
-    run_btn.click(start, inputs=user_input, outputs=output)
+    # 중요: gradio_client로 호출할 API 이름을 지정 (버튼 이벤트에 api_name).
+    # API 엔드포인트 이름은 "start_api"가 됨.
+    # (아래 추가로 동일 함수를 별도로 api 엔드포인트에 연결해도 됨.)
+    # 여기서는 클릭 핸들러에 api_name을 설정하려면 이렇게도 가능:
+    # run_btn_api.click(fn=run_all_api, inputs=[user_input], outputs=[api_output], api_name="start_api")
 
+# Launch - Space에서는 기본값으로 잘 동작함
 demo.launch()