m97j commited on
Commit
90bc37b
·
1 Parent(s): a18d920

Set config.py match with env variables such as HF_TOKEN

Browse files
Files changed (2) hide show
  1. config.py +17 -9
  2. model_loader.py +21 -5
config.py CHANGED
@@ -1,17 +1,25 @@
 
1
  import torch
 
2
 
3
- # 모델 경로
4
- BASE_MODEL = "Qwen/Qwen2.5-3B-Instruct"
5
- ADAPTER_MODEL = "m97j/npc-LoRA-fps"
 
 
 
6
 
7
  # 장치 설정
8
- DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
9
 
10
  # 토크나이저/모델 공통
11
- MAX_LENGTH = 1024
12
- NUM_FLAGS = 7 # flags.json 길이와 일치
13
 
14
  # 생성 파라미터
15
- GEN_MAX_NEW_TOKENS = 200
16
- GEN_TEMPERATURE = 0.7
17
- GEN_TOP_P = 0.9
 
 
 
 
1
+ import os
2
  import torch
3
+ from dotenv import load_dotenv
4
 
5
+ # .env 파일 로드 (로컬 개발 시)
6
+ load_dotenv()
7
+
8
+ # 모델 경로 (환경변수 없으면 기본값 사용)
9
+ BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen2.5-3B-Instruct")
10
+ ADAPTER_MODEL = os.getenv("ADAPTER_MODEL", "m97j/npc-LoRA-fps")
11
 
12
  # 장치 설정
13
+ DEVICE = os.getenv("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
14
 
15
  # 토크나이저/모델 공통
16
+ MAX_LENGTH = int(os.getenv("MAX_LENGTH", 1024))
17
+ NUM_FLAGS = int(os.getenv("NUM_FLAGS", 7)) # flags.json 길이와 일치
18
 
19
  # 생성 파라미터
20
+ GEN_MAX_NEW_TOKENS = int(os.getenv("GEN_MAX_NEW_TOKENS", 200))
21
+ GEN_TEMPERATURE = float(os.getenv("GEN_TEMPERATURE", 0.7))
22
+ GEN_TOP_P = float(os.getenv("GEN_TOP_P", 0.9))
23
+
24
+ # Hugging Face Token (Private 모델 접근용)
25
+ HF_TOKEN = os.getenv("HF_TOKEN")
model_loader.py CHANGED
@@ -2,13 +2,13 @@ import os, json, torch
2
  import torch.nn as nn
3
  from transformers import AutoTokenizer, AutoModelForCausalLM
4
  from peft import PeftModel
5
- from config import BASE_MODEL, ADAPTER_MODEL, DEVICE
6
 
7
  def get_current_branch():
8
  if os.path.exists("current_branch.txt"):
9
  with open("current_branch.txt", "r") as f:
10
  return f.read().strip()
11
- return "latest" # fallback
12
 
13
  class ModelWrapper:
14
  def __init__(self):
@@ -16,14 +16,30 @@ class ModelWrapper:
16
  self.flags_order = json.load(open(flags_path, encoding="utf-8"))["ALL_FLAGS"]
17
  self.num_flags = len(self.flags_order)
18
 
19
- self.tokenizer = AutoTokenizer.from_pretrained(ADAPTER_MODEL, use_fast=True)
 
 
 
 
 
20
  if self.tokenizer.pad_token is None:
21
  self.tokenizer.pad_token = self.tokenizer.eos_token
22
  self.tokenizer.padding_side = "right"
23
 
24
  branch = get_current_branch()
25
- base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto", trust_remote_code=True)
26
- self.model = PeftModel.from_pretrained(base, ADAPTER_MODEL, revision=branch, device_map="auto")
 
 
 
 
 
 
 
 
 
 
 
27
 
28
  hidden_size = self.model.config.hidden_size
29
  self.model.delta_head = nn.Linear(hidden_size, 2).to(DEVICE)
 
2
  import torch.nn as nn
3
  from transformers import AutoTokenizer, AutoModelForCausalLM
4
  from peft import PeftModel
5
+ from config import BASE_MODEL, ADAPTER_MODEL, DEVICE, HF_TOKEN
6
 
7
  def get_current_branch():
8
  if os.path.exists("current_branch.txt"):
9
  with open("current_branch.txt", "r") as f:
10
  return f.read().strip()
11
+ return "latest"
12
 
13
  class ModelWrapper:
14
  def __init__(self):
 
16
  self.flags_order = json.load(open(flags_path, encoding="utf-8"))["ALL_FLAGS"]
17
  self.num_flags = len(self.flags_order)
18
 
19
+ # 토큰 전달
20
+ self.tokenizer = AutoTokenizer.from_pretrained(
21
+ ADAPTER_MODEL,
22
+ use_fast=True,
23
+ token=HF_TOKEN
24
+ )
25
  if self.tokenizer.pad_token is None:
26
  self.tokenizer.pad_token = self.tokenizer.eos_token
27
  self.tokenizer.padding_side = "right"
28
 
29
  branch = get_current_branch()
30
+ base = AutoModelForCausalLM.from_pretrained(
31
+ BASE_MODEL,
32
+ device_map="auto",
33
+ trust_remote_code=True,
34
+ token=HF_TOKEN
35
+ )
36
+ self.model = PeftModel.from_pretrained(
37
+ base,
38
+ ADAPTER_MODEL,
39
+ revision=branch,
40
+ device_map="auto",
41
+ token=HF_TOKEN
42
+ )
43
 
44
  hidden_size = self.model.config.hidden_size
45
  self.model.delta_head = nn.Linear(hidden_size, 2).to(DEVICE)