Update api.py

api.py

@@ -1,12 +1,14 @@
-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      
+import requests
+import numpy as np
+import tensorflow as tf
+import asyncio
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import nltk
+nltk.download('punkt')
+from nltk.tokenize import 
     
 app = FastAPI()    
       
@@ -124,93 +126,127 @@ dummy_input = tf.zeros((1, max_len), dtype=tf.int32)  # 배치1, 시퀀스길이
 _ = 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):    
-    model_input = text_to_ids(f"<start> {prompt} <sep>")    
-    model_input = model_input[:max_len]    
-    generated = list(model_input)    
-    
-    tau = 5.0  # 초기 목표 surprise    
-    
-    for step in range(max_gen):    
-        pad_length = max(0, max_len - len(generated))    
-        input_padded = np.pad(generated, (0, pad_length), constant_values=pad_id)    
-        input_tensor = tf.convert_to_tensor([input_padded])    
-        logits = model(input_tensor, training=False)    
-        next_token_logits = logits[0, len(generated) - 1].numpy()    
-    
-        # 반복 페널티 적용    
-        token_counts = {}    
-        for t in generated:    
-            token_counts[t] = token_counts.get(t, 0) + 1    
-        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)    
-        cutoff = np.searchsorted(cumulative_probs, p, side='left') + 1    
-        filtered_indices = sorted_top_indices[:cutoff]    
-        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))    
-    
-        decoded_text = sp.decode(generated)    
-        # 특수 토큰 제거    
-        for token in ["<start>", "<sep>", "<end>"]:    
-            decoded_text = decoded_text.replace(token, "")    
-    
-        decoded_text = decoded_text.strip()    
-    
-        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", "안녕하세요")    
-        
-    # 스트리밍 응답으로 보냄    
+
+def extract_main_query(query):
+    words = query.split()
+    return " ".join(words[:3])
+
+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 summarize_text(text, top_n=3):
+    sentences = sent_tokenize(text)
+    if len(sentences) <= top_n:
+        return text
+    vectorizer = TfidfVectorizer(ngram_range=(1, 2), stop_words=['은', '는', '이', '가', '을', '를', '에', '에서'])
+    tfidf_matrix = vectorizer.fit_transform(sentences)
+    sim_matrix = cosine_similarity(tfidf_matrix, tfidf_matrix)
+    np.fill_diagonal(sim_matrix, 0)
+    scores = sim_matrix.sum(axis=1)
+    ranked_idx = np.argsort(scores)[::-1]
+    selected_idx = sorted(ranked_idx[:top_n])
+    summary = " ".join([sentences[i] for i in selected_idx])
+    return summary
+
+def simple_intent_classifier(text):
+    text = text.lower()
+    greet_keywords = ["안녕", "반가워", "이름", "누구", "소개", "어디서 왔", "정체", "몇 살", "너 뭐야"]
+    info_keywords = ["설명", "정보", "무엇", "뭐야", "어디", "누구", "왜", "어떻게", "종류", "개념"]
+    if any(kw in text for kw in greet_keywords):
+        return "인사"
+    elif any(kw in text for kw in info_keywords):
+        return "정보질문"
+    else:
+        return "일상대화"
+
+
+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):
+    model_input = text_to_ids(f"<start> {prompt} <sep>")
+    model_input = model_input[:max_len]
+    generated = list(model_input)
+    
+    tau = 5.0  # 초기 목표 surprise
+
+    for step in range(max_gen):
+        pad_length = max(0, max_len - len(generated))
+        input_padded = np.pad(generated, (0, pad_length), constant_values=pad_id)
+        input_tensor = tf.convert_to_tensor([input_padded])
+        logits = model(input_tensor, training=False)
+        next_token_logits = logits[0, len(generated) - 1].numpy()
+
+        # 반복 페널티 적용
+        token_counts = {}
+        for t in generated:
+            token_counts[t] = token_counts.get(t, 0) + 1
+        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()
+
+        sorted_indices = np.argsort(-probs)
+        top_indices = sorted_indices[:m]
+        top_probs = probs[top_indices]
+        top_probs /= top_probs.sum()
+
+        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)
+
+        sorted_top_indices = top_indices[np.argsort(-top_probs)]
+        sorted_top_probs = np.sort(top_probs)[::-1]
+        cumulative_probs = np.cumsum(sorted_top_probs)
+        cutoff = np.searchsorted(cumulative_probs, p, side='left') + 1
+        filtered_indices = sorted_top_indices[:cutoff]
+        filtered_probs = sorted_top_probs[:cutoff]
+        filtered_probs /= filtered_probs.sum()
+
+        final_token = np.random.choice(filtered_indices, p=filtered_probs)
+        generated.append(int(final_token))
+
+        decoded_text = decode_ids(generated)
+        for token in ["<start>", "<sep>", "<end>"]:
+            decoded_text = decoded_text.replace(token, "")
+        decoded_text = decoded_text.strip()
+
+        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):
+    intent = simple_intent_classifier(prompt)
+
+    if intent == "정보질문":
+        wiki_summary = get_wikipedia_summary(prompt)
+        summarized = summarize_text(wiki_summary, top_n=3)
+        yield f"『 \"{prompt}\" 에 대한 위키백과 요약입니다. 』\n\n{summarized}\n\n"
+
+    # 이후 일반 생성으로 이어감 (스트리밍)
+    gen = generate_text_mirostat_top_p(model, prompt)
+    for text_piece in gen:
+        yield text_piece
+        await asyncio.sleep(0.1)
+
+@app.get("/generate")
+async def generate(request: Request):
+    prompt = request.query_params.get("prompt", "안녕하세요")
     return StreamingResponse(async_generator_wrapper(prompt), media_type="text/plain")