Update api.py

api.py

@@ -1,139 +1,138 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
-import asyncio
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras import layers
-import sentencepiece as spm
-
-app = FastAPI()
-
-# SentencePiece 로드
-sp = spm.SentencePieceProcessor()
-sp.load("kolig_unigram.model")
-
-pad_id = sp.piece_to_id("<pad>")
-if pad_id == -1: pad_id = 0
-start_id = sp.piece_to_id("<start>")
-if start_id == -1: start_id = 1
-end_id = sp.piece_to_id("<end>")
-if end_id == -1: end_id = 2
-unk_id = sp.piece_to_id("<unk>")
-if unk_id == -1: unk_id = 3
-
-vocab_size = sp.get_piece_size()
-max_len = 100
-
-def text_to_ids(text):
-    return sp.encode(text, out_type=int)
-
-def ids_to_text(ids):
-    return sp.decode(ids)
-
-class RotaryPositionalEmbedding(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)
-
-    def call(self, x):
-        batch, heads, seq_len, depth = tf.unstack(tf.shape(x))
-        t = tf.range(seq_len, dtype=tf.float32)
-        freqs = tf.einsum('i,j->ij', t, self.inv_freq)
-        emb_sin = tf.sin(freqs)
-        emb_cos = tf.cos(freqs)
-        emb_cos = tf.reshape(emb_cos, [1, 1, seq_len, -1])
-        emb_sin = tf.reshape(emb_sin, [1, 1, seq_len, -1])
-        x1 = x[..., ::2]
-        x2 = x[..., 1::2]
-        x_rotated = tf.stack([
-            x1 * emb_cos - x2 * emb_sin,
-            x1 * emb_sin + x2 * emb_cos
-        ], axis=-1)
-        x_rotated = tf.reshape(x_rotated, tf.shape(x))
-        return x_rotated
-
-class SwiGLU(tf.keras.layers.Layer):
-    def __init__(self, d_model, d_ff):
-        super().__init__()
-        self.proj = tf.keras.layers.Dense(d_ff * 2)
-        self.out = tf.keras.layers.Dense(d_model)
-
-    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 GPTBlock(tf.keras.layers.Layer):
-    def __init__(self, d_model, d_ff, num_heads=8, dropout_rate=0.1, adapter_dim=64):
-        super().__init__()
-        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
-        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)
-        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)
-        self.rope = RotaryPositionalEmbedding(d_model // num_heads)
-
-    def call(self, x, training=False):
-        x_norm = self.ln1(x)
-        b, s, _ = tf.shape(x_norm)[0], tf.shape(x_norm)[1], tf.shape(x_norm)[2]
-        h = self.mha.num_heads
-        d = x_norm.shape[-1] // h
-
-        qkv = tf.reshape(x_norm, [b, s, h, d])
-        qkv = tf.transpose(qkv, [0, 2, 1, 3])
-        q = self.rope(qkv)
-        k = self.rope(qkv)
-        q = tf.reshape(tf.transpose(q, [0, 2, 1, 3]), [b, s, h * d])
-        k = tf.reshape(tf.transpose(k, [0, 2, 1, 3]), [b, s, h * d])
-
-        attn_out = self.mha(query=q, value=x_norm, key=k, use_causal_mask=True, training=training)
-        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 InteractGPT(tf.keras.Model):
-    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=8, dropout_rate=0.1):
-        super().__init__()
-        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model)
-        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]
-        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5)
-
-    def call(self, x, training=False):
-        x = self.token_embedding(x)
-        for block in self.blocks:
-            x = block(x, training=training)
-        x = self.ln_f(x)
-        logits = tf.matmul(x, self.token_embedding.embeddings, transpose_b=True)
-        return logits
-
-model = InteractGPT(vocab_size=vocab_size, seq_len=max_len, d_model=256, d_ff=1024, n_layers=6)
-
-dummy_input = tf.zeros((1, max_len), dtype=tf.int32)
-_ = model(dummy_input)
-model.load_weights("InteractGPT.weights.h5")
-print("모델 가중치 로드 완료!")
-
-
-async def generate_text_mirostat_top_p_with_buffer(model, prompt, max_len=100, max_gen=98,    
+from fastapi import FastAPI, Request    
+from fastapi.responses import StreamingResponse    
+import asyncio    
+import json        
+import numpy as np        
+import tensorflow as tf        
+from tensorflow.keras import layers        
+import sentencepiece as spm        
+import requests      
+    
+app = FastAPI()    
+      
+sp = spm.SentencePieceProcessor()      
+sp.load("kolig_unigram.model")      
+      
+pad_id = sp.piece_to_id("<pad>")      
+if pad_id == -1: pad_id = 0      
+start_id = sp.piece_to_id("<start>")      
+if start_id == -1: start_id = 1      
+end_id = sp.piece_to_id("< end >")      
+if end_id == -1: end_id = 2      
+unk_id = sp.piece_to_id("<unk>")      
+if unk_id == -1: unk_id = 3      
+      
+vocab_size = sp.get_piece_size()      
+max_len = 100      
+      
+def text_to_ids(text):      
+    return sp.encode(text, out_type=int)      
+      
+def ids_to_text(ids):      
+    return sp.decode(ids)      
+    
+class RotaryPositionalEmbedding(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)      
+      
+    def call(self, x):      
+        batch, heads, seq_len, depth = tf.unstack(tf.shape(x))      
+        t = tf.range(seq_len, dtype=tf.float32)      
+        freqs = tf.einsum('i,j->ij', t, self.inv_freq)      
+        emb_sin = tf.sin(freqs)      
+        emb_cos = tf.cos(freqs)      
+        emb_cos = tf.reshape(emb_cos, [1, 1, seq_len, -1])      
+        emb_sin = tf.reshape(emb_sin, [1, 1, seq_len, -1])      
+        x1 = x[..., ::2]      
+        x2 = x[..., 1::2]      
+        x_rotated = tf.stack([      
+            x1 * emb_cos - x2 * emb_sin,      
+            x1 * emb_sin + x2 * emb_cos      
+        ], axis=-1)      
+        x_rotated = tf.reshape(x_rotated, tf.shape(x))      
+        return x_rotated    
+    
+class SwiGLU(tf.keras.layers.Layer):    
+    def __init__(self, d_model, d_ff):    
+        super().__init__()    
+        self.proj = tf.keras.layers.Dense(d_ff * 2)    
+        self.out = tf.keras.layers.Dense(d_model)    
+    
+    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 GPTBlock(tf.keras.layers.Layer):    
+    def __init__(self, d_model, d_ff, num_heads=8, dropout_rate=0.1, adapter_dim=64):      
+        super().__init__()      
+        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
+        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)      
+        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)     
+        self.rope = RotaryPositionalEmbedding(d_model // num_heads)      
+      
+    def call(self, x, training=False):      
+        x_norm = self.ln1(x)      
+        b, s, _ = tf.shape(x_norm)[0], tf.shape(x_norm)[1], tf.shape(x_norm)[2]      
+        h = self.mha.num_heads      
+        d = x_norm.shape[-1] // h      
+      
+        qkv = tf.reshape(x_norm, [b, s, h, d])      
+        qkv = tf.transpose(qkv, [0, 2, 1, 3])      
+        q = self.rope(qkv)      
+        k = self.rope(qkv)      
+        q = tf.reshape(tf.transpose(q, [0, 2, 1, 3]), [b, s, h * d])      
+        k = tf.reshape(tf.transpose(k, [0, 2, 1, 3]), [b, s, h * d])      
+      
+        attn_out = self.mha(query=q, value=x_norm, key=k, use_causal_mask=True, training=training)      
+        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 InteractGPT(tf.keras.Model):      
+    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=8, dropout_rate=0.1):      
+        super().__init__()      
+        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model)      
+        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]      
+        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
+      
+    def call(self, x, training=False):      
+        x = self.token_embedding(x)      
+        for block in self.blocks:      
+            x = block(x, training=training)      
+        x = self.ln_f(x)      
+        logits = tf.matmul(x, self.token_embedding.embeddings, transpose_b=True)      
+        return logits      
+    
+model = InteractGPT(vocab_size=vocab_size, seq_len=max_len, d_model=256, d_ff=1024, n_layers=6)        
+      
+dummy_input = tf.zeros((1, max_len), dtype=tf.int32)  # 배치1, 시퀀스길이 max_len      
+_ = model(dummy_input)  # 모델이 빌드됨      
+model.load_weights("InteractGPT.weights.h5")      
+print("모델 가중치 로드 완료!")      
+    
+def generate_text_mirostat_top_p(model, prompt, max_len=100, max_gen=98,    
                                  temperature=1.0, min_len=20,    
-                                 repetition_penalty=1.2, eta=0.1, m=100, p=0.9, buffer_size=3):    
+                                 repetition_penalty=1.2, eta=0.1, m=100, p=0.9):    
     model_input = text_to_ids(f"<start> {prompt} <sep>")    
     model_input = model_input[:max_len]    
     generated = list(model_input)    
     
     tau = 5.0  # 초기 목표 surprise    
-    buffer_tokens = []    
     
     for step in range(max_gen):    
         pad_length = max(0, max_len - len(generated))    
@@ -149,26 +148,32 @@ async def generate_text_mirostat_top_p_with_buffer(model, prompt, max_len=100, m
         for token_id, count in token_counts.items():    
             next_token_logits[token_id] /= (repetition_penalty ** count)    
     
+        # 최소 길이 넘으면 종료 토큰 확률 낮추기    
         if len(generated) >= min_len:    
             next_token_logits[end_id] -= 5.0    
         next_token_logits[pad_id] -= 10.0    
     
+        # 온도 조절    
         next_token_logits = next_token_logits / temperature    
     
+        # --- 미로스타트 + Top-p 샘플링 ---    
         logits_stable = next_token_logits - np.max(next_token_logits)    
         probs = np.exp(logits_stable)    
         probs /= probs.sum()    
     
+        # 1. mirostat top-m 후보 추리기    
         sorted_indices = np.argsort(-probs)    
         top_indices = sorted_indices[:m]    
         top_probs = probs[top_indices]    
         top_probs /= top_probs.sum()    
     
+        # 2. mirostat 샘플링    
         sampled_index = np.random.choice(top_indices, p=top_probs)    
         sampled_prob = probs[sampled_index]    
         observed_surprise = -np.log(sampled_prob + 1e-9)    
         tau += eta * (observed_surprise - tau)    
     
+        # 3. top-p 필터링    
         sorted_top_indices = top_indices[np.argsort(-top_probs)]    
         sorted_top_probs = np.sort(top_probs)[::-1]    
         cumulative_probs = np.cumsum(sorted_top_probs)    
@@ -177,32 +182,35 @@ async def generate_text_mirostat_top_p_with_buffer(model, prompt, max_len=100, m
         filtered_probs = sorted_top_probs[:cutoff]    
         filtered_probs /= filtered_probs.sum()    
     
+        # 4. 최종 토큰 샘플링    
         final_token = np.random.choice(filtered_indices, p=filtered_probs)    
+        generated.append(int(final_token))    
     
-        if final_token == end_id:    
-            # 버퍼에 남은 거 다 출력    
-            if buffer_tokens:    
-                decoded = sp.decode(buffer_tokens)    
-                for token in ["<start>", "<sep>", "<end>"]:    
-                    decoded = decoded.replace(token, "")    
-                yield decoded.strip()    
-            break    
+        decoded_text = sp.decode(generated)    
+        # 특수 토큰 제거    
+        for token in ["<start>", "<sep>", "<end>"]:    
+            decoded_text = decoded_text.replace(token, "")    
     
-        if final_token in [start_id, pad_id] or sp.id_to_piece(final_token) == "<sep>":    
-            continue    
+        decoded_text = decoded_text.strip()    
     
-        generated.append(int(final_token))    
-        buffer_tokens.append(final_token)    
-    
-        if len(buffer_tokens) >= buffer_size or sp.id_to_piece(final_token).endswith("▁"):    
-            # 띄어쓰기 있는 토큰 나오거나 버퍼 꽉 찼으면 출력    
-            decoded = sp.decode(buffer_tokens)    
-            for token in ["<start>", "<sep>", "<end>"]:    
-                decoded = decoded.replace(token, "")    
-            yield decoded.strip()    
-            buffer_tokens = []
+        if len(generated) >= min_len and (final_token == end_id or decoded_text.endswith(('.', '!', '?'))):    
+            yield decoded_text    
+            break    
+            
+async def async_generator_wrapper(prompt: str):    
+    # 동기 제너레이터를 비동기로 감싸기    
+    loop = asyncio.get_event_loop()    
+    gen = generate_text_mirostat_top_p(model, prompt)    
+        
+    for text_piece in gen:    
+        yield text_piece    
+        # 토큰 생성 속도 조절 (0.1초 딜레이)    
+        await asyncio.sleep(0.1)    
+    
+@app.get("/generate")    
+async def generate(request: Request):    
+    # 쿼리 파라미터로 prompt 받음, 없으면 기본값    
+    prompt = request.query_params.get("prompt", "안녕하세요")    
         
-@app.get("/generate")
-async def generate(request: Request):
-    prompt = request.query_params.get("prompt", "안녕하세요")
-    return StreamingResponse(generate_text_mirostat_top_p_with_buffer(model, prompt), media_type="text/plain")
+    # 스트리밍 응답으로 보냄    
+    return StreamingResponse(async_generator_wrapper(prompt), media_type="text/plain")