Update api.py

api.py

@@ -228,7 +228,75 @@ def generate_text_sample(model, prompt, max_len=100, max_gen=98,
     for t in ["<start>", "<sep>", "<end>"]:
         decoded = decoded.replace(t, "")
     return decoded.strip()
-    
+
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.decomposition import TruncatedSVD
+from sklearn.metrics.pairwise import cosine_similarity
+
+class SimilarityMemory:
+    def __init__(self, n_components=100):
+        self.memory_texts = []
+        self.vectorizer = TfidfVectorizer()
+        self.svd = TruncatedSVD(n_components=n_components)
+        self.embeddings = None
+        self.fitted = False
+
+    def add(self, text: str):
+        self.memory_texts.append(text)
+        self._update_embeddings()
+
+    def _update_embeddings(self):
+        # 텍스트가 1개 이상일 때만 벡터화 및 차원 축소 진행
+        if len(self.memory_texts) == 0:
+            self.embeddings = None
+            self.fitted = False
+            return
+
+        # 벡터화
+        X = self.vectorizer.fit_transform(self.memory_texts)
+
+        # 차원 축소
+        n_comp = min(self.svd.n_components, X.shape[1], len(self.memory_texts)-1)
+        if n_comp <= 0:
+            # 축소할 차원이 없으면 그냥 TF-IDF 벡터 사용
+            self.embeddings = X.toarray()
+            self.fitted = True
+            return
+
+        self.svd = TruncatedSVD(n_components=n_comp)
+        self.embeddings = self.svd.fit_transform(X)
+        self.fitted = True
+
+    def retrieve(self, query: str, top_k=3):
+        if not self.fitted or self.embeddings is None or len(self.memory_texts) == 0:
+            return []
+
+        # 쿼리 벡터화 + 차원 축소 (fit_transform이 아니라 transform 해야 함)
+        Xq = self.vectorizer.transform([query])
+        if self.svd.n_components > Xq.shape[1] or self.svd.n_components > len(self.memory_texts) - 1:
+            q_emb = Xq.toarray()
+        else:
+            q_emb = self.svd.transform(Xq)
+
+        # 코사인 유사도 계산
+        sims = cosine_similarity(q_emb, self.embeddings)[0]
+        top_indices = sims.argsort()[::-1][:top_k]
+
+        return [self.memory_texts[i] for i in top_indices]
+
+# 테스트
+memory = SimilarityMemory()
+
+memory.add("어제는 기분이 별로였어")
+memory.add("영화 보러 갔다가 친구랑 싸웠어")
+memory.add("카페에서 공부했는데 집중이 잘 됐어")
+
+def merge_prompt_with_memory(prompt: str, memories: list):
+    context = "\n".join(f"과거: {mem}" for mem in memories)
+    return f"{context} {prompt}"
+
+
 def mismatch_tone(input_text, output_text):  
     if "ㅋㅋ" in input_text and not re.search(r'ㅋㅋ|ㅎ|재밌|놀|만나|맛집|여행', output_text):  
         return True  
@@ -328,7 +396,6 @@ def parse_math_question(text):
     except:
         return "계산할 수 없는 수식이에요. 다시 한번 확인해 주세요!"
 
-# 최종 응답 함수
 def respond(input_text):
     intent = simple_intent_classifier(input_text)
 
@@ -348,10 +415,13 @@ def respond(input_text):
         summary = summarize_from_wikipedia(keyword)
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
-    # 일상 대화: 샘플링 + fallback
-    response = generate_text_sample(model, input_text)
+    # ✅ 기억 기반 병합 추가
+    related_memories = memory.retrieve(input_text, top_k=3)
+    merged_prompt = merge_prompt_with_memory(input_text, related_memories)
+
+    response = generate_text_sample(model, merged_prompt)
     if not is_valid_response(response) or mismatch_tone(input_text, response):
-        response = generate_text_sample(model, input_text)
+        response = generate_text_sample(model, merged_prompt)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)