Update api.py

api.py

@@ -229,7 +229,6 @@ def generate_text_sample(model, prompt, max_len=100, max_gen=98,
         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
@@ -238,9 +237,9 @@ class SimilarityMemory:
     def __init__(self, n_components=100):
         self.memory_texts = []
         self.vectorizer = TfidfVectorizer()
-        self.svd = None
+        self.svd = TruncatedSVD(n_components=n_components)
         self.embeddings = None
-        self.n_components = n_components
+        self.fitted = False
 
     def add(self, text: str):
         self.memory_texts.append(text)
@@ -249,42 +248,44 @@ class SimilarityMemory:
     def _update_embeddings(self):
         if len(self.memory_texts) == 0:
             self.embeddings = None
-            self.svd = None
+            self.fitted = False
             return
 
         X = self.vectorizer.fit_transform(self.memory_texts)
-
-        n_comp = min(self.n_components, X.shape[1], len(self.memory_texts) - 1)
+        n_comp = min(self.svd.n_components, X.shape[1], len(self.memory_texts)-1)
         if n_comp <= 0:
             self.embeddings = X.toarray()
-            self.svd = None
+            self.fitted = True
             return
 
-        svd = TruncatedSVD(n_components=n_comp)
-        self.embeddings = svd.fit_transform(X)
-        self.svd = svd
+        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 self.embeddings is None:
+        if not self.fitted or self.embeddings is None or len(self.memory_texts) == 0:
             return []
 
         Xq = self.vectorizer.transform([query])
-        if self.svd is not None:
-            q_emb = self.svd.transform(Xq)
-        else:
+        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()
+    def process_input(self, new_text: str, top_k=3):
+        """자동으로 기억 저장하고, 유사한 기억 찾아서 합친 프롬프트 생성"""
+        related_memories = self.retrieve(new_text, top_k=top_k)
+        self.add(new_text)
+        return self.merge_prompt(new_text, related_memories)
 
-def merge_prompt_with_memory(prompt: str, memories: list):
-    context = "\n".join(memories)
-    return f"{context}\n{prompt}"
+    def merge_prompt(self, prompt: str, memories: list):
+        context = "\n".join(memories)
+        return f"{context}\n\n{prompt}" if context else prompt
 
 
 def mismatch_tone(input_text, output_text):