Initial commit

app.py

@@ -8,33 +8,28 @@ from models.model_loader import load_emotion_model, load_fallback_model, load_em
 from schemas import AskReq, AskRes
 from pathlib import Path
 from rag.rag_generator import set_embedder
-
-# 모델 이름
-EMOTION_MODEL_NAME = "tae898/emoberta-base-ko"
-FALLBACK_MODEL_NAME = "skt/ko-gpt-trinity-1.2B-v0.5"
-EMBEDDER_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
-
-# 절대 경로 기준 모델 디렉토리 설정
-BASE_DIR = Path(__file__).resolve().parent  # ai_server/
-EMOTION_MODEL_DIR = Path(os.getenv("EMOTION_MODEL_DIR", BASE_DIR / "models" / "emotion-classification-model"))
-FALLBACK_MODEL_DIR = Path(os.getenv("FALLBACK_MODEL_DIR", BASE_DIR / "models" / "fallback-npc-model"))
-EMBEDDER_MODEL_DIR = Path(os.getenv("EMBEDDER_MODEL_DIR", BASE_DIR / "models" / "sentence-embedder"))
+from config import (
+    EMOTION_MODEL_NAME, EMOTION_MODEL_DIR,
+    FALLBACK_MODEL_NAME, FALLBACK_MODEL_DIR,
+    EMBEDDER_MODEL_NAME, EMBEDDER_MODEL_DIR,
+    HF_TOKEN, BASE_DIR
+)
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Emotion
-    emo_tokenizer, emo_model = load_emotion_model(EMOTION_MODEL_NAME, EMOTION_MODEL_DIR)
+    emo_tokenizer, emo_model = load_emotion_model(EMOTION_MODEL_NAME, EMOTION_MODEL_DIR, token=HF_TOKEN)
     app.state.emotion_tokenizer = emo_tokenizer
     app.state.emotion_model = emo_model
 
     # Fallback
-    fb_tokenizer, fb_model = load_fallback_model(FALLBACK_MODEL_NAME, FALLBACK_MODEL_DIR)
+    fb_tokenizer, fb_model = load_fallback_model(FALLBACK_MODEL_NAME, FALLBACK_MODEL_DIR, token=HF_TOKEN)
     app.state.fallback_tokenizer = fb_tokenizer
     app.state.fallback_model = fb_model
 
     # Embedder
-    embedder = load_embedder(EMBEDDER_MODEL_NAME, EMBEDDER_MODEL_DIR)
+    embedder = load_embedder(EMBEDDER_MODEL_NAME, EMBEDDER_MODEL_DIR, token=HF_TOKEN)
     app.state.embedder = embedder
     set_embedder(embedder)  # 추가
 

config.py

@@ -1,32 +1,17 @@
+from pathlib import Path
 import os
 
-# Hugging Face Spaces serve URL (공개 설정이므로 직접 주소 사용 가능)
-HF_SERVE_URL = os.getenv(
-    "HF_SERVE_URL",
-    "https://m97j-PersonaChatEngine.hf.space"
-)
+BASE_DIR = Path(__file__).resolve().parent  # ai_server/
 
-# 요청 타임아웃 (초 단위)
-HF_TIMEOUT = float(os.getenv("HF_TIMEOUT", "25"))
+# Hugging Face Token
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# RAG 항상 사용 (토글이 아니라 고정 사용)
-RAG_ENABLED = True
+# 모델 이름
+EMOTION_MODEL_NAME = "tae898/emoberta-base-ko"
+FALLBACK_MODEL_NAME = "skt/ko-gpt-trinity-1.2B-v0.5"
+EMBEDDER_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
 
-# 생성 파라미터 기본값 (요청마다 override 가능)
-GENERATION_CONFIG = {
-    "max_new_tokens": int(os.getenv("GEN_MAX_NEW_TOKENS", "220")),
-    "temperature": float(os.getenv("GEN_TEMPERATURE", "0.7")),
-    "top_p": float(os.getenv("GEN_TOP_P", "0.9")),
-    "repetition_penalty": float(os.getenv("GEN_REPETITION_PENALTY", "1.1")),
-    "do_sample": True
-}
-
-'''
-# 모델 정보 (추후 확장 가능)
-MODEL_INFO = {
-    "base_model": "meta-llama/Meta-Llama-3-8B",
-    "adapter": "m97j/PersonaAdapter-v1",
-    "serve_mode": "hf_spaces",  # 또는 "local", "api"
-}
-
-'''
+# 모델 디렉토리
+EMOTION_MODEL_DIR = Path(os.getenv("EMOTION_MODEL_DIR", BASE_DIR / "models" / "emotion-classification-model"))
+FALLBACK_MODEL_DIR = Path(os.getenv("FALLBACK_MODEL_DIR", BASE_DIR / "models" / "fallback-npc-model"))
+EMBEDDER_MODEL_DIR = Path(os.getenv("EMBEDDER_MODEL_DIR", BASE_DIR / "models" / "sentence-embedder"))

models/download_models.py

@@ -1,19 +0,0 @@
-import os
-from transformers import AutoTokenizer, AutoModel
-
-# Hugging Face token from secret file
-with open("/run/secrets/HF_TOKEN", "r") as f:
-    token = f.read().strip()
-
-models = {
-    'emotion': ('tae898/emoberta-base-ko', './models/emotion-classification-model'),
-    'fallback': ('skt/ko-gpt-trinity-1.2B-v0.5', './models/fallback-npc-model'),
-    'embedder': ('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2', './models/sentence-embedder')
-}
-
-for _, (_, path) in models.items():
-    os.makedirs(path, exist_ok=True)
-
-for name, path in models.values():
-    AutoModel.from_pretrained(name, cache_dir=path, use_auth_token=token)
-    AutoTokenizer.from_pretrained(name, cache_dir=path, use_auth_token=token)

models/model_loader.py

@@ -7,10 +7,10 @@ from transformers import (
 from sentence_transformers import SentenceTransformer
 
 
-def load_emotion_model(model_name: str, model_dir: Path):
+def load_emotion_model(model_name: str, model_dir: Path, token: str = None):
     if not model_dir.exists() or not any(model_dir.iterdir()):
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_auth_token=token)
+        model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, use_auth_token=token)
         tokenizer.save_pretrained(model_dir)
         model.save_pretrained(model_dir)
 
@@ -19,10 +19,10 @@ def load_emotion_model(model_name: str, model_dir: Path):
     return tokenizer, model
 
 
-def load_fallback_model(model_name: str, model_dir: Path):
+def load_fallback_model(model_name: str, model_dir: Path, token: str = None):
     if not model_dir.exists() or not any(model_dir.iterdir()):
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_auth_token=token)
+        model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, use_auth_token=token)
         tokenizer.save_pretrained(model_dir)
         model.save_pretrained(model_dir)
 
@@ -31,9 +31,9 @@ def load_fallback_model(model_name: str, model_dir: Path):
     return tokenizer, model
 
 
-def load_embedder(model_name: str, model_dir: Path):
+def load_embedder(model_name: str, model_dir: Path, token: str = None):
     if not model_dir.exists() or not any(model_dir.iterdir()):
-        embedder = SentenceTransformer(model_name)
+        embedder = SentenceTransformer(model_name, use_auth_token=token)
         embedder.save(str(model_dir))
 
     embedder = SentenceTransformer(str(model_dir))