Upload app.py

app.py

@@ -1,40 +1,32 @@
 import os
 import json
-import threading
 import torch
-import gradio as gr
 from huggingface_hub import login
 from sentence_transformers import SentenceTransformer
 import faiss
-
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer
-)
+import numpy as np
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from peft import PeftModel
-
+import threading
 
 # ============================================================
-# 0. 환경 설정
+# 1. 환경 설정 + 로그인
 # ============================================================
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-if HF_TOKEN:
-    login(token=HF_TOKEN)
+HF_TOKEN = os.getenv("HF_TOKEN")  # ← secret variable에서 불러옴
+login(token=HF_TOKEN)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
+print("Device:", device)
 
 # ============================================================
-# 1. 경로 설정
+# 2. 경로 설정
 # ============================================================
 BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
-LORA_DIR = "peft_lora"
-DOC_PATH = "rule.json"
-
+LORA_DIR = "./peft_lora"           # 서버 경로
+DOC_PATH = "./rule.json"           # 문서 파일
 
 # ============================================================
-# 2. RAG 문서 로드
+# 3. RAG 문서 로드 + FAISS 구축
 # ============================================================
 with open(DOC_PATH, "r", encoding="utf-8") as f:
     documents = json.load(f)
@@ -46,23 +38,19 @@ embedding_model = SentenceTransformer(
     device=device
 )
 
-import numpy as np
 doc_embs = embedding_model.encode(
-    doc_texts,
-    convert_to_numpy=True,
-    show_progress_bar=True
+    doc_texts, convert_to_numpy=True
 ).astype("float32")
 
 dim = doc_embs.shape[1]
 index = faiss.IndexFlatL2(dim)
 index.add(doc_embs)
 
+print("FAISS index built:", index.ntotal)
 
 # ============================================================
-# 3. LLM + LoRA 로드
+# 4. LLM + LoRA 로드
 # ============================================================
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
-
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
@@ -70,27 +58,30 @@ model = AutoModelForCausalLM.from_pretrained(
     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
+
 model = PeftModel.from_pretrained(
     model,
     LORA_DIR,
     torch_dtype=torch.float16,
-    device_map="auto"
+    device_map="auto",
 )
 
+model = model.to(device)
 model.eval()
 
-
 # ============================================================
-# 4. RAG 검색
+# 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]]
 
-
 # ============================================================
-# 5. 프롬프트 생성
+# 6. 프롬프트 생성
 # ============================================================
 def build_prompt(persona, instruction, query, retrieved_docs):
     context = "\n".join([f"- {d['text']}" for d in retrieved_docs])
@@ -110,21 +101,25 @@ def build_prompt(persona, instruction, query, retrieved_docs):
 ### 답변:
 """
 
-
 # ============================================================
-# 6. Streaming LLM (End of Answer 이전까지만 출력)
+# 7. Streaming Chat
 # ============================================================
-def stream_generate(prompt, max_new_tokens=256):
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
+def stream_chat(persona, instruction, user_query, max_new_tokens=256):
+
+    retrieved = retrieve(user_query, k=3)
+    prompt = build_prompt(persona, instruction, user_query, retrieved)
 
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-    thread = threading.Thread(
-        target=lambda:
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    END_TOKENS = [
+        "End of Answer", "### 검토 결과:", "### 최종 답변",
+        "※", ">", "**답변**", "---", "###", "**"
+    ]
+
+    def run_gen():
+        with torch.no_grad():
             model.generate(
                 **inputs,
                 max_new_tokens=max_new_tokens,
@@ -134,51 +129,56 @@ def stream_generate(prompt, max_new_tokens=256):
                 repetition_penalty=1.2,
                 streamer=streamer
             )
-    )
-    thread.start()
 
-    partial_text = ""
-
-    for token in streamer:
-        partial_text += token
-
-        # End of Answer 기준으로 스트림 중단
-        if "End of Answer" in partial_text:
-            partial_text = partial_text.split("End of Answer")[0]
-            yield partial_text.strip()
-            return
+    thread = threading.Thread(target=run_gen)
+    thread.start()
 
-        yield partial_text
+    full = ""
+    for tok in streamer:
+        print(tok, end="", flush=True)
+        full += tok
+        for e in END_TOKENS:
+            if e in full:
+                print()
+                return
 
+    print()
 
 # ============================================================
-# 7. Gradio 인터페이스 함수
+# 8. 페르소나 목록
 # ============================================================
-def gradio_reply(persona, instruction, query):
-    retrieved = retrieve(query, k=3)
-    prompt = build_prompt(persona, instruction, query, retrieved)
-    return stream_generate(prompt)
-
+persona_group = [
+    ("당신은 원칙을 지키되 상황에 따라 유연하게 판단하는 시각을 가지고 있다...", "박세연"),
+    ("당신은 공정한 규칙과 원칙을 중시하면서, 개인의 성과와 능력을 인정해 차등...", "김창준"),
+    ("규율과 자율의 균형을 지키며, 능력과 성과를 기준으로 판단한다...", "이상기"),
+    ("규율을 기반으로 하지만 유연하며, 분배는 중립적이고 개선을 추구한다...", "채훈"),
+    ("자율을 존중하되 최소한의 규율을 유지하며, 기여도와 개선을 균형 있게 반영...", "용우"),
+    ("규율과 공정을 기반으로 안정적인 운영을 추구하며, 균등·개선·친목 간의 균형...", "형진")
+]
 
 # ============================================================
-# 8. Gradio UI
+# 9. 프로그램 실행 (입력 받는 부분)
 # ============================================================
-with gr.Blocks() as demo:
-
-    gr.Markdown("KORMo-10B + LoRA + RAG Streaming Demo (End-of-Answer Truncated)")
-
-    persona = gr.Textbox(label="페르소나")
-    instruction = gr.Textbox(label="규칙/지침")
-    query = gr.Textbox(label="질문")
+if __name__ == "__main__":
 
-    output = gr.Textbox(label="응답", lines=12)
+    query = input("질문을 입력하세요: ")
 
-    btn = gr.Button("Generate")
+    instruction = """
+    당신은 해당 페르소나의 성격을 가진 심판관입니다.
+    반드시 3문장만 말하십시오.
+    각 문장은 30자 이내.
+    규정을 우선하여 답하세요.
+    판단 근거 포함.
+    반복 금지.
+    """
 
-    btn.click(
-        fn=gradio_reply,
-        inputs=[persona, instruction, query],
-        outputs=output
-    )
+    for persona_text, persona_name in persona_group:
+        print("\n====================")
+        print(f"### {persona_name} ###")
+        print("====================")
 
-demo.launch()
+        stream_chat(
+            persona=persona_text,
+            instruction=instruction,
+            user_query=query
+        )