Update Test.py

Test.py

@@ -1,362 +1,237 @@
-!pip install sentencepiece
+import os, json, random, numpy as np, tensorflow as tf
+from tensorflow.keras import layers, Model
 import sentencepiece as spm
-
-# 불러오기
-import os, json, numpy as np, tensorflow as tf
 import requests
-print('1')
-
-tf.get_logger().setLevel("ERROR")
-SEED = 42
-tf.random.set_seed(SEED)
-np.random.seed(SEED)
-
-# TPU 초기화
-try:
-    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="local")
-    tf.tpu.experimental.initialize_tpu_system(resolver)
-    strategy = tf.distribute.TPUStrategy(resolver)
-    print("✅ TPU 초기화 완료:", resolver.cluster_spec().as_dict())
-    on_tpu = True
-except Exception as e:
-    print("⚠️ TPU 미사용, GPU/CPU로 진행:", e)
-    strategy = tf.distribute.get_strategy()
-    on_tpu = False
-
-# Mixed precision
-from tensorflow.keras import mixed_precision
-import tensorflow as tf
-from tensorflow.keras import layers, activations, initializers
-policy = mixed_precision.Policy("mixed_bfloat16" if on_tpu else "float32")
-mixed_precision.set_global_policy(policy)
-print("✅ Mixed precision:", policy)
-
-# =======================
-# 1) 파일 다운로드
-# =======================
-def download_file(url, save_path):
-    r = requests.get(url, stream=True)
-    r.raise_for_status()
-    with open(save_path, "wb") as f:
-        for chunk in r.iter_content(8192):
-            f.write(chunk)
-    print(f"✅ {save_path} 저장됨")
-
-DATA_PATH = "converted.jsonl"
-TOKENIZER_PATH = "ko_unigram.model"
-
-if not os.path.exists(DATA_PATH):
-    download_file(
-        "https://huggingface.co/datasets/Yuchan5386/SFT/resolve/main/data_shuffled_1.jsonl?download=true",
-        DATA_PATH
-    )
-
-if not os.path.exists(TOKENIZER_PATH):
-    download_file(
-        "https://huggingface.co/Yuchan5386/inlam-70m-instruct/resolve/main/unigram.model?download=true",
-        TOKENIZER_PATH
-    )
 
-sp = spm.SentencePieceProcessor(TOKENIZER_PATH)
+# ===============================
+# 0️⃣ 환경 설정
+# ===============================
+TOKENIZER_PATH = "bpe.model"
+DATA_PATH = "corpus.txt"  # 36M 문장 텍스트 파일
+MAX_LEN = 128
+EMBED_DIM = 384
+LATENT_DIM = 384
+BATCH_SIZE = 400
+NEGATIVE_RATIO = 1  # negative sample 수
 
+def download_file(url, save_path):
+    if not os.path.exists(save_path):
+        print(f"Downloading {save_path} ...")
+        r = requests.get(url, stream=True)
+        r.raise_for_status()
+        with open(save_path, "wb") as f:
+            for chunk in r.iter_content(8192*2):
+                f.write(chunk)
+        print(f"✅ {save_path} saved")
+
+download_file("https://huggingface.co/datasets/OpenLab-NLP/ko-corpus/resolve/main/bpe.model?download=true", TOKENIZER_PATH)
+download_file("https://huggingface.co/datasets/OpenLab-NLP/ko-corpus/resolve/main/shuffled_corpus%20(1).txt?download=true", DATA_PATH)
+
+# ===============================
+# 2️⃣ 토크나이저 준비
+# ===============================
+sp = spm.SentencePieceProcessor(TOKENIZER_PATH)
 pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
-start_id = sp.piece_to_id("<start>")
-sep_id = sp.piece_to_id("<sep>")
-end_id = sp.piece_to_id("<end>")
-unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
-print(f"✅ Vocabulary size: {vocab_size}")
-
-max_len = 1024
-batch_size = 128
-
-def text_to_ids(text):
-    return sp.encode(text, out_type=int)
-def ids_to_text(ids):
-    return sp.decode(ids)
-
-def jsonl_stream(file_path):
-    with open(file_path, "r", encoding="utf-8") as f:
-        for line in f:
-            data = json.loads(line)
-            conversations = data.get("conversations", [])
-            for i in range(0, len(conversations) - 1, 2):
-                human_msg = conversations[i]
-                gpt_msg   = conversations[i + 1]
-                if human_msg.get("from") != "human" or gpt_msg.get("from") != "gpt":
-                    continue
-                prompt   = human_msg.get("value", "").strip()
-                response = gpt_msg.get("value", "").strip()
-                full = f"<start> {prompt} <sep> {response} <end>"
-                if "<sep>" not in full:
-                    continue
-                sep_index  = full.index("<sep>")
-                input_text = full[:sep_index + len("<sep>")].strip()
-                target_text = full[sep_index + len("<sep>"):].strip()
-
-                input_ids  = text_to_ids(input_text)
-                target_ids = text_to_ids(target_text + " <end>")
-
-                available_len = max_len - len(input_ids)
-                if available_len <= 0:
-                    input_ids = input_ids[-max_len:]
-                    target_ids = []
-                    target_mask = [0] * len(input_ids)
-                else:
-                    target_ids = target_ids[:available_len]
-                    target_mask = [0] * len(input_ids) + [1] * len(target_ids)
-
-                full_input = input_ids + target_ids
-                pad_len = max_len - len(full_input)
-                full_input += [pad_id] * pad_len
-                target_mask += [0] * pad_len
-
-                target_seq = full_input[1:] + [end_id]
-                target_seq = target_seq[:max_len]
-
-                masked_target = [
-                    t if m == 1 else pad_id
-                    for t, m in zip(target_seq, target_mask)
-                ]
-
-                yield (
-                    tf.convert_to_tensor(full_input, dtype=tf.int32),
-                    tf.convert_to_tensor(masked_target, dtype=tf.int32)
-                )
+
+def encode_sentence(sentence, max_len=MAX_LEN):
+    return sp.encode(sentence, out_type=int)[:max_len]
+
+def pad_sentence(tokens):
+    return tokens + [pad_id]*(MAX_LEN - len(tokens))
+
+def gen_pairs_streaming(txt_path=DATA_PATH, negative_ratio=NEGATIVE_RATIO):
+    with open(txt_path, "r", encoding="utf-8") as f:
+        sentences = [line.strip() for line in f if line.strip()]
+    while True:
+        for s1 in sentences:
+            # positive pair (자기 자신)
+            x1 = pad_sentence(encode_sentence(s1))
+            yield (x1, x1), 1.0
+
+            # negative pairs (자기 자신 제외)
+            for _ in range(negative_ratio):
+                s2 = s1
+                while s2 == s1:
+                    s2 = random.choice(sentences)
+                x2 = pad_sentence(encode_sentence(s2))
+                yield (x1, x2), 0.0
 
 dataset = tf.data.Dataset.from_generator(
-    lambda: jsonl_stream(DATA_PATH),
-    output_signature=(
-        tf.TensorSpec(shape=(max_len,), dtype=tf.int32),
-        tf.TensorSpec(shape=(max_len,), dtype=tf.int32),
-    ),
-)
-dataset = dataset.shuffle(1000, seed=SEED).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)
-
-with strategy.scope():
-    dist_dataset = strategy.experimental_distribute_dataset(dataset)
-    
-class RotaryPositionalEmbedding(tf.keras.layers.Layer):
-    def __init__(self, dim):  
-        super().__init__()  
-        inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))  
-        self.inv_freq = tf.constant(inv_freq, dtype=tf.float32)
+    lambda: gen_pairs_streaming(),
+    output_types=((tf.int32, tf.int32), tf.float32),
+    output_shapes=(((MAX_LEN,), (MAX_LEN,)), ())
+).shuffle(1024).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
 
-    def call(self, x):
-        b, h, s, d = tf.unstack(tf.shape(x))
-        t = tf.range(s, dtype=tf.float32)
-        freqs = tf.einsum('i,j->ij', t, self.inv_freq)
-        dtype = x.dtype
-        emb_sin = tf.cast(tf.sin(freqs), dtype)
-        emb_cos = tf.cast(tf.cos(freqs), dtype)
-        emb_cos = tf.reshape(emb_cos, [1,1,s,-1])
-        emb_sin = tf.reshape(emb_sin, [1,1,s,-1])
-        x1, x2 = x[..., ::2], x[..., 1::2]
-        x_rot = tf.stack([x1*emb_cos - x2*emb_sin, x1*emb_sin + x2*emb_cos], axis=-1)
-        x_rot = tf.reshape(x_rot, tf.shape(x))
-        return x_rot
-
-class SwiGLU(tf.keras.layers.Layer):
-    def __init__(self, d_model, d_ff):
+class EncoderBlock(tf.keras.layers.Layer):
+    def __init__(self, embed_dim=EMBED_DIM, ff_dim=1152, seq_len=MAX_LEN):
         super().__init__()
-        self.proj = tf.keras.layers.Dense(d_ff)
-        self.out  = tf.keras.layers.Dense(d_model)
+        self.fc1 = layers.Dense(ff_dim)
+        self.fc2 = layers.Dense(embed_dim)
+        self.fc3 = layers.Dense(ff_dim)
+        self.fc4 = layers.Dense(embed_dim)
+
+        self.w_proj = self.add_weight(
+            shape=(embed_dim, embed_dim),
+            initializer="glorot_uniform",
+            trainable=True
+        )
+
+        self.alpha2 = layers.Dense(1)
+
+        self.ln = layers.LayerNormalization(epsilon=1e-5)
+        self.ln1 = layers.LayerNormalization(epsilon=1e-5)
+        self.ln2 = layers.LayerNormalization(epsilon=1e-5)
+        
     def call(self, x):
-        x_proj = self.proj(x)
-        x_val, x_gate = tf.split(x_proj, 2, axis=-1)
-        return self.out(x_val * tf.nn.silu(x_gate))
-
-class FlashAttentionMHA(layers.Layer):
-    def __init__(self, d_model, num_heads=8, dropout_rate=0.1):
-        super().__init__()
-        self.d_model = d_model
-        self.num_heads = num_heads
-        self.dh = d_model // num_heads
-
-        self.q_proj = layers.Dense(d_model, use_bias=False)
-        self.k_proj = layers.Dense(d_model, use_bias=False)
-        self.v_proj = layers.Dense(d_model, use_bias=False)
-        self.out_proj = layers.Dense(d_model, use_bias=False)
-        self.dropout = layers.Dropout(dropout_rate)
-        self.rope = RotaryPositionalEmbedding(self.dh)
-
-    @tf.function(jit_compile=True)
-    def call(self, x, training=False, causal=False):
-        B, N, D = tf.shape(x)[0], tf.shape(x)[1], x.shape[2]
-
-        # Q,K,V: (B, N, num_heads, dh)
-        Q = tf.reshape(self.q_proj(x), [B, N, self.num_heads, self.dh])
-        K = tf.reshape(self.k_proj(x), [B, N, self.num_heads, self.dh])
-        V = tf.reshape(self.v_proj(x), [B, N, self.num_heads, self.dh])
-
-        # transpose for attention: (B, num_heads, N, dh)
-        Q = tf.transpose(Q, [0,2,1,3])
-        K = tf.transpose(K, [0,2,1,3])
-        V = tf.transpose(V, [0,2,1,3])
-
-        # ROPE 적용
-        Q = self.rope(Q)
-        K = self.rope(K)
-
-        # Scaled dot-product
-        scale = tf.cast(self.dh ** -0.5, x.dtype)
-        Q = Q * scale
-        attn_scores = tf.matmul(Q, K, transpose_b=True)
-
-        if causal:
-            mask = tf.linalg.band_part(tf.ones((N,N), dtype=x.dtype), -1, 0)
-            attn_scores = attn_scores * mask - 1e9 * (1 - mask)
-
-        attn_weights = tf.nn.softmax(attn_scores, axis=-1)
-        attn_weights = self.dropout(attn_weights, training=training)
-        out = tf.matmul(attn_weights, V)  # (B, h, N, dh)
-        out = tf.transpose(out, [0,2,1,3])
-        out = tf.reshape(out, [B, N, D])
-        out = self.out_proj(out)
-        return out
-
-
-class GPTBlock(tf.keras.layers.Layer):
-    def __init__(self, d_model, d_ff, num_heads=12, dropout_rate=0.1, adapter_dim=64):
-        super().__init__()
-        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
-        self.mha = FlashAttentionMHA(d_model, num_heads, dropout_rate=dropout_rate)
-        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)
-        self.adapter_down = tf.keras.layers.Dense(adapter_dim, activation='gelu')
-        self.adapter_up   = tf.keras.layers.Dense(d_model)
-        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
-        self.ffn = SwiGLU(d_model, d_ff)
-        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)
-
-    def call(self, x, training=False):
-        x_norm = self.ln1(x)
-        attn_out = self.mha(x_norm, training=training, causal=True)
-        attn_out = self.dropout1(attn_out, training=training)
-        adapter_out = self.adapter_up(self.adapter_down(attn_out))
-        attn_out = attn_out + adapter_out
-        x = x + attn_out
-        ffn_out = self.ffn(self.ln2(x))
-        x = x + self.dropout2(ffn_out, training=training)
-        return x
-
-class InLaM(tf.keras.Model):
-    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=12, dropout_rate=0.1):
+        x_norm = self.ln(x)
+        x = self.fc1(x_norm)
+        g, v = tf.split(x, 2, axis=-1)
+        x = tf.nn.silu(g) * v
+        x = self.fc2(x)
+
+        x = tf.matmul(x, x, transpose_b=True)          # (B,L,L)
+        x = tf.tensordot(x, self.w_proj, axes=[-1, 0]) # (B,L,D)
+
+        v = tf.nn.softmax(self.alpha2(v), axis=1) * x
+        x_norm = x_norm + self.ln2(v)
+
+        x = self.fc3(x_norm)
+        g, v = tf.split(x, 2, axis=-1)
+        x = tf.nn.silu(g) * v
+        x = self.fc4(x)
+
+        return x_norm + self.ln1(x)
+
+
+class L2NormLayer(layers.Layer):
+    def __init__(self, axis=1, epsilon=1e-10, **kwargs):
+        super().__init__(**kwargs)
+        self.axis = axis
+        self.epsilon = epsilon
+    def call(self, inputs):
+        return tf.math.l2_normalize(inputs, axis=self.axis, epsilon=self.epsilon)
+    def get_config(self):
+        return {"axis": self.axis, "epsilon": self.epsilon, **super().get_config()}
+
+class SentenceEncoder(tf.keras.Model):
+    def __init__(self, vocab_size, embed_dim=384, latent_dim=384, max_len=128, pad_id=pad_id):
         super().__init__()
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        
-        # Embedding 레이어 (bfloat16)
-        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model, dtype="bfloat16")
-        
-        # Transformer Blocks
-        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]
-        
-        # Final LayerNorm
-        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5, dtype="bfloat16")
-    def call(self, x, training=False):
-        # Embedding
-        x = self.token_embedding(x)  # (batch, seq_len, d_model)
-        for block in self.blocks:
-            x = block(x, training=training)
+        self.pad_id = pad_id
+        self.embed = layers.Embedding(vocab_size, embed_dim)
+        self.pos_embed = layers.Embedding(input_dim=max_len, output_dim=embed_dim)
+        self.blocks = [EncoderBlock() for _ in range(1)]
+        self.attn_pool = layers.Dense(1)
+        self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
+        self.latent = layers.Dense(latent_dim, activation=None)  # tanh 제거
+        self.l2norm = L2NormLayer()  # 추가
 
-        x = self.ln_f(x)  # (batch, seq_len, d_model)
-        embed_weights = self.token_embedding.weights[0]  # (vocab_size, d_model)
-        logits = tf.matmul(x, embed_weights, transpose_b=True)  # (batch, seq_len, vocab_size)
-        
-        # float32로 캐스팅 (손실 계산 등에서 안정성 확보)
-        return tf.cast(logits, tf.float32)
-
-# =======================
-# 손실/메트릭 정의
-# =======================
-def smoothed_loss_keras(y_true, y_pred, eps=0.1):
-    y_true = tf.cast(y_true, tf.int32)
-    mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-    vocab = tf.shape(y_pred)[-1]
-    y_true_oh = tf.one_hot(y_true, depth=vocab, dtype=tf.float32)
-    y_true_ls = (1.0 - eps) * y_true_oh + eps / tf.cast(vocab, tf.float32)
-    log_probs = tf.nn.log_softmax(y_pred, axis=-1)
-    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1)
-    per_tok = per_tok * mask
-    return tf.reduce_sum(per_tok) / (tf.reduce_sum(mask) + 1e-8)
-
-def masked_accuracy(y_true, y_pred):
-    y_true = tf.cast(y_true, tf.int32)
-    mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-    pred_id = tf.argmax(y_pred, axis=-1, output_type=tf.int32)
-    acc = tf.cast(tf.equal(y_true, pred_id), tf.float32) * mask
-    return tf.reduce_sum(acc) / (tf.reduce_sum(mask) + 1e-8)
-
-def masked_perplexity(y_true, y_pred, eps=0.1):
-    y_true = tf.cast(y_true, tf.int32)
-    mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-    vocab = tf.shape(y_pred)[-1]
-    y_true_oh = tf.one_hot(y_true, depth=vocab, dtype=tf.float32)
-    y_true_ls = (1.0 - eps) * y_true_oh + eps / tf.cast(vocab, tf.float32)
-    log_probs = tf.nn.log_softmax(y_pred, axis=-1)
-    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1)
-    per_tok = per_tok * mask
-    mean_loss = tf.reduce_sum(per_tok) / (tf.reduce_sum(mask) + 1e-8)
-    return tf.exp(mean_loss)
-
-
-# =======================
-# 모델 생성 & 컴파일
-# =======================
-with strategy.scope():
-    model = InLaM(vocab_size=vocab_size, seq_len=max_len, d_model=768, d_ff=768*4, n_layers=12)
-    dummy_input = tf.zeros((batch_size, max_len), dtype=tf.int32)
-    _ = model(dummy_input, training=False)
-    model.summary()
-
-    optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.9, beta_2=0.95, epsilon=1e-8, clipnorm=1.0)
-    model.compile(optimizer=optimizer, loss=smoothed_loss_keras, metrics=[masked_accuracy, masked_perplexity])
-
-    # 학습
-    history = model.fit(dist_dataset, epochs=1, verbose=1)
-
-# =======================
-# 가중치 저장
-# =======================
-model.save_weights("tf_model.weights.h5")
-print("✅ 모델 가중치 저장 완료!")
-
-# =======================
-# 샘플 생성 함수
-# =======================
-def generate_text_topp(model, prompt, max_len=115, max_gen=98, p=0.9, temperature=0.68, min_len=20):
-    model_input = text_to_ids(f"<start> {prompt} <sep>")
-    model_input = model_input[:max_len]
-    generated = list(model_input)
-    
-    for step in range(max_gen):
-        input_seq = generated[-max_len:] if len(generated) > max_len else generated
-        input_padded = np.pad(input_seq, (0, max_len - len(input_seq)), constant_values=pad_id)
-        input_tensor = tf.convert_to_tensor([input_padded], dtype=tf.int32)
-        
-        logits = model(input_tensor, training=False).numpy()[0, len(input_seq)-1]
-        logits[end_id] -= 5.0
-        logits[pad_id] -= 10.0
-        
-        probs = tf.nn.softmax(logits / temperature).numpy()
-        sorted_idx = np.argsort(probs)[::-1]
-        sorted_probs = probs[sorted_idx]
-        cumulative = np.cumsum(sorted_probs)
-        cutoff = np.searchsorted(cumulative, p)
-        top_idx = sorted_idx[:cutoff + 1]
-        top_probs = sorted_probs[:cutoff + 1] / sorted_probs[:cutoff + 1].sum()
-        
-        next_token = int(np.random.choice(top_idx, p=top_probs))
-        if next_token == end_id and len(generated) >= min_len:
+    def call(self, x):
+        positions = tf.range(tf.shape(x)[1])[tf.newaxis, :]
+        x_embed = self.embed(x) + self.pos_embed(positions)
+        mask = tf.cast(tf.not_equal(x, self.pad_id), tf.float32)
+        x = x_embed
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+
+        scores = self.attn_pool(x)
+        scores = tf.where(tf.equal(mask[..., tf.newaxis], 0), -1e9, scores)
+        scores = tf.nn.softmax(scores, axis=1)
+        pooled = tf.reduce_sum(x * scores, axis=1)
+
+        latent = self.latent(pooled)
+        return self.l2norm(latent)  # L2 정규화 후 반환
+
+# ===============================
+# 5️⃣ Cosine similarity layer + Contrastive Loss
+# ===============================
+class CosineSimilarityLayer(layers.Layer):
+    def call(self, inputs):
+        v1, v2 = inputs
+        return tf.reduce_sum(v1 * v2, axis=-1)  # 이미 L2 정규화돼서 dot product = cosine similarity
+
+def contrastive_loss(margin=0.5):
+    def loss(y_true, y_pred):
+        y_true = tf.cast(y_true, tf.float32)
+        dist = 1 - y_pred
+        pos_loss = y_true * tf.square(dist)
+        neg_loss = (1 - y_true) * tf.square(tf.maximum(margin - dist, 0))
+        return tf.reduce_mean(pos_loss + neg_loss)
+    return loss
+
+encoder = SentenceEncoder(vocab_size=vocab_size)
+
+# ===============================
+# 6️⃣ 시암 모델 정의
+# ===============================
+input1 = tf.keras.Input(shape=(MAX_LEN,), dtype=tf.int32)
+input2 = tf.keras.Input(shape=(MAX_LEN,), dtype=tf.int32)
+v1 = encoder(input1)
+v2 = encoder(input2)
+cos_sim = CosineSimilarityLayer()([v1, v2])
+siamese_model = tf.keras.Model([input1, input2], cos_sim)
+siamese_model.compile(optimizer=tf.keras.optimizers.Adam(1e-5), loss=contrastive_loss(margin=0.5))
+siamese_model.summary()
+# ===============================
+# 7️⃣ 학습
+# ===============================
+#steps_per_epoch = 36757266 // 400
+steps_per_epoch = 1000000 // 400
+# generator 기반 streaming 학습
+siamese_model.fit(dataset, epochs=1, steps_per_epoch=steps_per_epoch)  # steps_per_epoch는 필요에 따라 조절
+encoder.save_weights("encoder.weights.h5")
+siamese_model.save_weights("siamese_model.weights.h5")
+
+# ===============================
+# 8️⃣ corpus 벡터 생성 + 캐싱 (안전하게 새로 생성)
+# ===============================
+LIMIT = 1000  # 검색용 corpus 문장 수
+prompts = []
+
+# prompts 먼저 읽기
+with open(DATA_PATH, "r", encoding="utf-8") as f:
+    for i, line in enumerate(f):
+        if i >= LIMIT:
             break
-        generated.append(next_token)
-    
-    return ids_to_text(generated)
-
-# =======================
-# 테스트 생성
-# =======================
-prompt = "안녕하세요! 한국 밴드에 대해 궁금한 것이 있어요!"
-sample_text = generate_text_topp(model, prompt, p=0.9)
-print("\n===== 생성 결과 =====\n")
-print(sample_text)
+        line = line.strip()
+        if line:
+            prompts.append(line)
+
+def get_sentence_vector(sentence):
+    tokens = pad_sentence(encode_sentence(sentence))
+    return encoder(np.array([tokens])).numpy()[0]
+
+# corpus_vectors 항상 새로 생성 (기존 npy 무시)
+corpus_vectors = np.stack([get_sentence_vector(p) for p in prompts]).astype(np.float16)
+np.save("corpus_vectors.npy", corpus_vectors)
+
+# norms 계산
+corpus_norms = np.linalg.norm(corpus_vectors, axis=1)
+
+# ===============================
+# 9️⃣ 검색 함수
+# ===============================
+def search(query, top_k=3):
+    q_vec = get_sentence_vector(query).astype(np.float16)
+    sims = corpus_vectors @ q_vec
+    sims /= (corpus_norms * np.linalg.norm(q_vec) + 1e-8)
+
+    # top_k 안전 처리
+    top_k = min(top_k, len(prompts))
+    top_idx = np.argsort(sims)[::-1][:top_k]
+
+    return [(prompts[i], float(sims[i])) for i in top_idx]
+
+# ===============================
+# 🔟 테스트
+# ===============================
+query = "우리가 핸드폰, 배를 세계에서 제일 잘 만드는 것 이상으로 사랑을 제일 잘 실천할 수 있는 능력, 자질, 저력이 우리에게 있다."
+results = search(query)
+for p, s in results:
+    print(f"Prompt: {p}\n유사도: {s:.3f}\n---")
+
+query = "안녕하세요! 오늘 날씨 어떤가요?"
+results = search(query)
+for p, s in results:
+    print(f"Prompt: {p}\n유사도: {s:.3f}\n---")