api.py

import requests
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers  
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import sentencepiece as spm
import re
app = FastAPI()    


from fastapi.middleware.cors import CORSMiddleware

origins = [
    "https://insect5386.github.io",
    "https://insect5386.github.io/insect5386"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
      
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_typical(model, prompt, max_len=100, max_gen=98,
                          temperature=0.50, min_len=20,
                          repetition_penalty=1.2, typical_p=0.80):
    
    def typical_filtering(logits, typical_p):
        probs = np.exp(logits - np.max(logits))
        probs /= probs.sum()

        log_probs = np.log(probs + 1e-9)
        entropy = -np.sum(probs * log_probs)

        shifted = np.abs(-log_probs - entropy)
        sorted_idx = np.argsort(shifted)
        sorted_probs = probs[sorted_idx]

        cum_probs = np.cumsum(sorted_probs)
        cutoff = np.searchsorted(cum_probs, typical_p) + 1

        final_idx = sorted_idx[:cutoff]
        final_probs = probs[final_idx]
        final_probs /= final_probs.sum()

        return final_idx, final_probs

    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):
        pad_len = max(0, max_len - len(generated))
        input_padded = np.pad(generated, (0, pad_len), constant_values=pad_id)
        input_tensor = tf.convert_to_tensor([input_padded])

        logits = model(input_tensor, training=False)
        next_logits = logits[0, len(generated) - 1].numpy()

        # 반복 억제
        for t in set(generated):
            count = generated.count(t)
            next_logits[t] /= (repetition_penalty ** count)

        # 조기 종료 방지
        if len(generated) < min_len:
            next_logits[end_id] -= 5.0
        next_logits[pad_id] -= 10.0

        # 온도 적용
        next_logits = next_logits / temperature

        # Typical Sampling 적용
        final_idx, final_probs = typical_filtering(next_logits, typical_p=typical_p)
        sampled = np.random.choice(final_idx, p=final_probs)

        generated.append(int(sampled))

        decoded = sp.decode(generated)
        for t in ["<start>", "<sep>", "<end>"]:
            decoded = decoded.replace(t, "")
        decoded = decoded.strip()

        if len(generated) >= min_len and (sampled == end_id or decoded.endswith(('.', '!', '?'))):
            return decoded  # ← yield 대신 return

def is_valid_response(response):
    if len(response.strip()) < 2:
        return False
    if re.search(r'[ㄱ-ㅎㅏ-ㅣ]{3,}', response):
        return False
    if len(response.split()) < 2:
        return False
    if response.count(' ') < 2:
        return False
    if any(tok in response.lower() for tok in ['hello', 'this', 'ㅋㅋ']):
        return False
    return True

def extract_main_query(text):
    sentences = re.split(r'[.?!]\s*', text)
    sentences = [s.strip() for s in sentences if s.strip()]
    if not sentences:
        return text
    last = sentences[-1]
    last = re.sub(r'[^가-힣a-zA-Z0-9 ]', '', last)
    particles = ['이', '가', '은', '는', '을', '를', '의', '에서', '에게', '한테', '보다']
    for p in particles:
        last = re.sub(rf'\b(\w+){p}\b', r'\1', last)
    return last.strip()

def get_wikipedia_summary(query):
    cleaned_query = extract_main_query(query)
    url = f"https://ko.wikipedia.org/api/rest_v1/page/summary/{cleaned_query}"
    res = requests.get(url)
    if res.status_code == 200:
        return res.json().get("extract", "요약 정보를 찾을 수 없습니다.")
    else:
        return "위키백과에서 정보를 가져올 수 없습니다."

def simple_intent_classifier(text):
    text = text.lower()
    greet_keywords = ["안녕", "반가워", "이름", "누구", "소개", "어디서 왔", "정체", "몇 살", "너 뭐야"]
    info_keywords = ["설명", "정보", "무엇", "뭐야", "어디", "누구", "왜", "어떻게", "종류", "개념"]
    math_keywords = ["더하기", "빼기", "곱하기", "나누기", "루트", "제곱", "+", "-", "*", "/", "=", "^", "√", "계산", "몇이야", "얼마야"]

    if any(kw in text for kw in greet_keywords):
        return "인사"
    elif any(kw in text for kw in info_keywords):
        return "정보질문"
    elif any(kw in text for kw in math_keywords):
        return "수학질문"
    else:
        return "일상대화"

def parse_math_question(text):
    text = text.replace("곱하기", "*").replace("더하기", "+").replace("빼기", "-").replace("나누기", "/").replace("제곱", "*2")
    text = re.sub(r'루트\s(\d+)', r'math.sqrt(\1)', text)
    try:
        result = eval(text)
        return f"정답은 {result}입니다."
    except:
        return "계산할 수 없는 수식이에요. 다시 한번 확인해 주세요!"

# 전체 응답 함수
def respond(input_text):
    intent = simple_intent_classifier(input_text)

    if "이름" in input_text:
        return "제 이름은 InteractGPT입니다."

    if "누구" in input_text:
        return "저는 InteractGPT이라고 해요."

    if intent == "수학질문":
        return parse_math_question(input_text)

    if intent == "인사":
        return "반가워요! 무엇을 도와드릴까요?"
        
    if intent == "정보질문":
        keyword = re.sub(r"(에 대해|에 대한|에 대해서)?\s*(설명해줘|알려줘|뭐야|개념|정의|정보)?", "", input_text).strip()
        if not keyword:
            return "어떤 주제에 대해 궁금한가요?"
        summary = get_wikipedia_summary(keyword)
        return f"{summary}\n다른 궁금한 점 있으신가요?"

    return generate_text_typical(model, input_text)

async def async_generator_wrapper(prompt: str):
    gen = generate_text_typical(model, prompt)
    for text_piece in gen:
        yield text_piece
        await asyncio.sleep(0.1)


from fastapi.responses import PlainTextResponse

@app.get("/generate", response_class=PlainTextResponse)
async def generate(request: Request):
    prompt = request.query_params.get("prompt", "안녕하세요")
    response_text = respond(prompt)
    return response_text