1# @title HybriKo-117M (Exp6 - Step 1962) Final
2import os
3import sys
4import yaml
5import torch
6import torch.nn.functional as F
7import sentencepiece as spm
8from huggingface_hub import hf_hub_download, list_repo_files
9
10# 1. 의존성 설치 및 설정
11!pip install transformers sentencepiece pyyaml huggingface_hub -q
12
13# 설정
14REPO_ID = "Yaongi/hybridko-exp6"
15CODE_REPO_ID = "Yaongi/HybriKo-117M"
16TARGET_CHECKPOINT = "checkpoint_step_1962.pt"
17DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
18print(f"🖥️ 사용 디바이스: {DEVICE}")
19
20# 2. 모델 코드 다운로드
21try:
22 print(f"📥 모델 코드 다운로드 시도 (from {REPO_ID})...")
23 config_py_path = hf_hub_download(REPO_ID, "configuration_hybridko.py")
24 model_py_path = hf_hub_download(REPO_ID, "modeling_hybridko.py")
25except Exception:
26 print(f"⚠️ 코드 파일이 없어 {CODE_REPO_ID}에서 다운로드합니다.")
27 config_py_path = hf_hub_download(CODE_REPO_ID, "configuration_hybridko.py")
28 model_py_path = hf_hub_download(CODE_REPO_ID, "modeling_hybridko.py")
29
30sys.path.insert(0, os.path.dirname(config_py_path))
31try:
32 from configuration_hybridko import HybriKoConfig
33 from modeling_hybridko import HybriKoModel
34except ImportError:
35 # Colab Import 오류 방지
36 import importlib
37 spec = importlib.util.spec_from_file_location("configuration_hybridko", config_py_path)
38 configuration_hybridko = importlib.util.module_from_spec(spec)
39 spec.loader.exec_module(configuration_hybridko)
40 HybriKoConfig = configuration_hybridko.HybriKoConfig
41 spec = importlib.util.spec_from_file_location("modeling_hybridko", model_py_path)
42 modeling_hybridko = importlib.util.module_from_spec(spec)
43 spec.loader.exec_module(modeling_hybridko)
44 HybriKoModel = modeling_hybridko.HybriKoModel
45
46# 3. 모델 및 체크포인트 로드
47try:
48 config_path = hf_hub_download(REPO_ID, "config.yaml")
49 with open(config_path, 'r') as f:
50 config_data = yaml.safe_load(f)
51 config = HybriKoConfig(**config_data)
52except:
53 config = HybriKoConfig()
54
55model = HybriKoModel(config)
56
57try:
58 # 1. checkpoints 폴더 2. 루트 순서 탐색
59 try:
60 ckpt_path = hf_hub_download(REPO_ID, f"checkpoints/{TARGET_CHECKPOINT}")
61 except:
62 ckpt_path = hf_hub_download(REPO_ID, TARGET_CHECKPOINT)
63
64 checkpoint = torch.load(ckpt_path, map_location=DEVICE)
65 state_dict = checkpoint.get("model_state_dict", checkpoint)
66 model.load_state_dict(state_dict)
67 print(f"✅ 모델 로드 완료: {TARGET_CHECKPOINT}")
68except Exception as e:
69 print(f"❌ 모델 로드 실패: {e}")
70
71model.to(DEVICE)
72model.eval()
73
74# 4. 토크나이저 로드
75try:
76 try:
77 tokenizer_path = hf_hub_download(REPO_ID, "tokenizer/HybriKo_tok.model")
78 except:
79 tokenizer_path = hf_hub_download(REPO_ID, "HybriKo_tok.model")
80 sp = spm.SentencePieceProcessor()
81 sp.Load(tokenizer_path)
82except Exception as e:
83 print(f"❌ 토크나이저 로드 실패: {e}")
84 sp = None
85
86# 5. 고급 생성 함수 (Repetition Penalty 지원)
87def generate(
88 text,
89 max_len=128,
90 temp=0.6, # 최적값: 0.6
91 top_k=50,
92 top_p=0.9, # 최적값: 0.9
93 repetition_penalty=1.15 # 최적값: 1.15
94):
95 if not sp or not text: return ""
96
97 input_ids = torch.tensor([[2] + sp.EncodeAsIds(text)]).to(DEVICE)
98 with torch.no_grad():
99 for _ in range(max_len):
100 idx = input_ids[:, -model.config.max_seq_len:]
101 outputs = model(idx)
102 logits = outputs.logits[:, -1] / temp
103
104 # 1. Repetition Penalty
105 if repetition_penalty > 1.0:
106 for i in range(input_ids.shape[1]):
107 token_id = input_ids[0, i]
108 if logits[0, token_id] < 0:
109 logits[0, token_id] *= repetition_penalty
110 else:
111 logits[0, token_id] /= repetition_penalty
112
113 # 2. Top-K
114 if top_k is not None:
115 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
116 logits[logits < v[:, [-1]]] = float("-inf")
117
118 # 3. Top-P
119 if top_p is not None:
120 sorted_logits, sorted_indices = torch.sort(logits, descending=True)
121 cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
122 sorted_indices_to_remove = cumulative_probs > top_p
123 sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
124 sorted_indices_to_remove[:, 0] = 0
125 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
126 logits[indices_to_remove] = float("-inf")
127
128 probs = F.softmax(logits, dim=-1)
129 next_token = torch.multinomial(probs, 1)
130
131 if next_token.item() == 3: # EOS Token
132 break
133
134 input_ids = torch.cat([input_ids, next_token], dim=1)
135
136 decoded = sp.DecodeIds(input_ids[0].tolist())
137 return decoded
138
139print("\n" + "="*30)
140print(f"🤖 HybriKo 채팅 시작 (종료: 'q')")
141print(f"⚙️ 옵션: Temp={0.6}, Rep_Pen={1.15}")
142print("="*30)
143
144while True:
145 try:
146 user_input = input("\nUser: ")
147 if user_input.lower() in ['q', 'quit', 'exit']:
148 break
149
150 response = generate(user_input)
151 print(f"HybriKo: {response}")
152 except KeyboardInterrupt:
153 break
154 except Exception as e:
155 print(f"Error: {e}")
156
