Update app.py

app.py

@@ -1,85 +1,188 @@
+import os
 import torch
+import json
+import threading
+import time
+import faiss
+import numpy as np
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# ======================
-# 1. 모델 로딩
-# ======================
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"   # 원하는 모델로 변경 가능
+from sentence_transformers import SentenceTransformer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+)
+from peft import PeftModel
+from huggingface_hub import login
+
+# ============================================================
+# 1. 로그인 (Colab 전용 코드 제거, Space 환경 변수 사용)
+# ============================================================
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if 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 = "./kormo_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 구축
+# ============================================================
+with open(DOC_PATH, "r", encoding="utf-8") as f:
+    documents = json.load(f)
+
+doc_texts = [d["text"] for d in documents]
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+embedding_model = SentenceTransformer("jhgan/ko-sroberta-multitask", device=device)
 
+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.add(doc_embs)
+
+print("FAISS index built. Total docs =", index.ntotal)
+
+# ============================================================
+# 4. LLM + LoRA 로드
+# ============================================================
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,      # 양자화 아님
-    device_map="auto"
+    BASE_MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True
 )
 
-# ======================
-# 2. 스트리밍 함수
-# ======================
-def generate_stream(prompt):
-    """
-    이 함수는 Gradio에서 스트리밍이 되도록
-    yield 로 토큰 단위 출력하는 제너레이터입니다.
-    """
+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",
+)
+model = model.to(device)
+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. 프롬프트 구성 (네 코드 그대로)
+# ============================================================
+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, k=3, max_new_tokens=256):
+
+    retrieved = retrieve(user_query, k=k)
+    prompt = build_prompt(persona, instruction, user_query, retrieved)
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-    # 스트리밍 생성기 호출
-    streamer = tokenizer.decode
-
-    # 모델의 generate()는 streaming 지원 X → 아래는 수동 스트리밍 구현
-    # (HuggingFace TextStreamer 대신 토큰 단위 수동 처리)
-    
-    generated = inputs["input_ids"]
-    past_key_values = None
-
-    for _ in range(512):  # 최대 생성 토큰 수
-        outputs = model(
-            input_ids=generated if past_key_values is None else generated[:, -1:],
-            past_key_values=past_key_values,
-            use_cache=True
-        )
-        logits = outputs.logits[:, -1, :]
-        past_key_values = outputs.past_key_values
-
-        next_token = torch.argmax(logits, dim=-1, keepdim=True)
-        generated = torch.cat([generated, next_token], dim=-1)
-
-        # 디코드 후 사용자에게 스트리밍 전달
-        text = tokenizer.decode(generated[0], skip_special_tokens=True)
-        yield text
-
-        # 종료 토큰 발견 시 stop
-        if next_token.item() in tokenizer.eos_token_id if isinstance(tokenizer.eos_token_id, list) else [tokenizer.eos_token_id]:
-            break
-
-
-# ======================
-# 3. Gradio UI
-# ======================
-with gr.Blocks() as demo:
-    gr.Markdown("# 🚀 Custom LLM Streaming Demo (No Quantization)")
-
-    with gr.Row():
-        input_box = gr.Textbox(label="Prompt 입력", lines=4)
-    
-    output_box = gr.Textbox(label="응답 (Streaming)")
-
-    generate_button = gr.Button("생성")
-
-    # 버튼 클릭시 스트리밍 연결
-    generate_button.click(
-        fn=generate_stream,
-        inputs=input_box,
-        outputs=output_box
+    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)
+    thread.start()
+
+    full = ""
+    for token in streamer:
+        full += token
+        yield token
+
+# ============================================================
+# 8. Gradio UI
+# ============================================================
+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
     )
 
-# ======================
-# 4. 실행
-# ======================
-if __name__ == "__main__":
-    demo.launch()
+app.launch()