Update app.py

app.py

@@ -1,184 +1,211 @@
-import os
-import json
-import torch
-from huggingface_hub import login
-from sentence_transformers import SentenceTransformer
-import faiss
-import numpy as np
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from peft import PeftModel
-import threading
-
-# ============================================================
-# 1. 환경 설정 + 로그인
-# ============================================================
-HF_TOKEN = os.getenv("HF_TOKEN")  # ← secret variable에서 불러옴
-login(token=HF_TOKEN)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print("Device:", device)
-
-# ============================================================
-# 2. 경로 설정
-# ============================================================
-BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
-LORA_DIR = "./peft_lora"           # 서버 경로
-DOC_PATH = "./rule.json"           # 문서 파일
-
-# ============================================================
-# 3. 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")
-
-dim = doc_embs.shape[1]
-index = faiss.IndexFlatL2(dim)
-index.add(doc_embs)
-
-print("FAISS index built:", index.ntotal)
-
-# ============================================================
-# 4. LLM + LoRA 로드
-# ============================================================
-model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    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",
-)
-
-model = model.to(device)
-model.eval()
-
-# ============================================================
-# 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. 프롬프트 생성
-# ============================================================
-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}
-
-### 답변:
-"""
-
-# ============================================================
-# 7. Streaming Chat
-# ============================================================
-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)
-
-    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,
-                do_sample=True,
-                top_p=0.9,
-                temperature=0.7,
-                repetition_penalty=1.2,
-                streamer=streamer
-            )
-
-    thread = threading.Thread(target=run_gen)
-    thread.start()
-
-    full = ""
-    for tok in streamer:
-        print(tok, end="", flush=True)
-        full += tok
-        for e in END_TOKENS:
-            if e in full:
-                print()
-                return
-
-    print()
-
-# ============================================================
-# 8. 페르소나 목록
-# ============================================================
-persona_group = [
-    ("당신은 원칙을 지키되 상황에 따라 유연하게 판단하는 시각을 가지고 있다...", "박세연"),
-    ("당신은 공정한 규칙과 원칙을 중시하면서, 개인의 성과와 능력을 인정해 차등...", "김창준"),
-    ("규율과 자율의 균형을 지키며, 능력과 성과를 기준으로 판단한다...", "이상기"),
-    ("규율을 기반으로 하지만 유연하며, 분배는 중립적이고 개선을 추구한다...", "채훈"),
-    ("자율을 존중하되 최소한의 규율을 유지하며, 기여도와 개선을 균형 있게 반영...", "용우"),
-    ("규율과 공정을 기반으로 안정적인 운영을 추구하며, 균등·개선·친목 간의 균형...", "형진")
-]
-
-# ============================================================
-# 9. 프로그램 실행 (입력 받는 부분)
-# ============================================================
-if __name__ == "__main__":
-
-    query = input("질문을 입력하세요: ")
-
-    instruction = """
-    당신은 해당 페르소나의 성격을 가진 심판관입니다.
-    반드시 3문장만 말하십시오.
-    각 문장은 30자 이내.
-    규정을 우선하여 답하세요.
-    판단 근거 포함.
-    반복 금지.
-    """
-
-    for persona_text, persona_name in persona_group:
-        print("\n====================")
-        print(f"### {persona_name} ###")
-        print("====================")
-
-        stream_chat(
-            persona=persona_text,
-            instruction=instruction,
-            user_query=query
-        )
+import os
+import json
+import torch
+import threading
+import numpy as np
+import faiss
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
+
+from huggingface_hub import login
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from peft import PeftModel
+
+
+# ============================================================
+# FastAPI 설정
+# ============================================================
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ============================================================
+# 1. 환경 설정 + 로그인
+# ============================================================
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("환경 변수 HF_TOKEN이 설정되지 않음")
+
+login(token=HF_TOKEN)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Device:", device)
+
+# ============================================================
+# 2. 경로 설정
+# ============================================================
+BASE_MODEL = "KORMo-Team/KORMo-10B-sft"
+LORA_DIR = "./peft_lora"
+DOC_PATH = "./rule.json"
+
+
+# ============================================================
+# 3. 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")
+
+dim = doc_embs.shape[1]
+index = faiss.IndexFlatL2(dim)
+index.add(doc_embs)
+
+print("FAISS index built:", index.ntotal)
+
+
+# ============================================================
+# 4. LLM + LoRA 로드
+# ============================================================
+model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto",
+    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 if device == "cuda" else torch.float32,
+    device_map="auto",
+)
+
+model.eval()
+
+
+# ============================================================
+# 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. 프롬프트 생성
+# ============================================================
+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}
+
+### 답변:
+"""
+
+
+# ============================================================
+# 7. Streaming Chat (동기적 모아서 반환)
+# ============================================================
+def run_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)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    def 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,
+                streamer=streamer
+            )
+
+    thread = threading.Thread(target=generate)
+    thread.start()
+
+    full_result = ""
+    for token in streamer:
+        full_result += token
+
+    return full_result.strip()
+
+
+# ============================================================
+# 8. 페르소나 목록
+# ============================================================
+persona_group = [
+    ("당신은 원칙을 지키되 상황에 따라 유연하게 판단하는 시각을 가지고 있다...", "박세연"),
+    ("당신은 공정한 규칙과 원칙을 중시하면서, 개인의 성과와 능력을 인정해 차등...", "김창준"),
+    ("규율과 자율의 균형을 지키며, 능력과 성과를 기준으로 판단한다...", "이상기"),
+    ("규율을 기반으로 하지만 유연하며, 분배는 중립적이고 개선을 추구한다...", "채훈"),
+    ("자율을 존중하되 최소한의 규율을 유지하며, 기여도와 개선을 균형 있게 반영...", "용우"),
+    ("규율과 공정을 기반으로 안정적인 운영을 추구하며, 균등·개선·친목 간의 균형...", "형진")
+]
+
+
+# ============================================================
+# 9. API 입력 모델
+# ============================================================
+class UserQuery(BaseModel):
+    query: str
+
+
+# ============================================================
+# 10. 실제 API 라우트
+# ============================================================
+@app.post("/ask")
+async def ask_api(payload: UserQuery):
+    user_query = payload.query
+
+    instruction = """
+    당신은 해당 페르소나의 성격을 가진 심판관입니다.
+    반드시 3문장만 말하십시오.
+    각 문장은 30자 이내.
+    규정을 우선하여 답하세요.
+    판단 근거 포함.
+    반복 금지.
+    """
+
+    results = {}
+
+    for persona_text, persona_name in persona_group:
+        answer = run_chat(persona_text, instruction, user_query)
+        results[persona_name] = answer
+
+    return {"query": user_query, "answers": results}
+
+
+# ============================================================
+# Root Index
+# ============================================================
+@app.get("/")
+async def root():
+    return {"status": "running", "message": "KORMo + LoRA + RAG Persona API"}