Update app.py

app.py

@@ -6,11 +6,29 @@ from datasets import load_dataset
 
 app = FastAPI()
 
-# 1. 임베딩 모델 로드
+# 1. MiniLM 모델 로드 (데이터 벡터화)
 embed_model = SentenceTransformer("all-MiniLM-L6-v2")
 
-# 2. Hugging Face에서 MedRAG 데이터셋 로드
+# 2. MedRAG 데이터셋 로드
 dataset = load_dataset("MedRAG/textbooks", split="train", streaming=True)
 
 # 3. 데이터 변환
 texts = [entry["content"] for entry in dataset]  # "content" 필드 활용
+
+# 4. 벡터 임베딩 생성 및 FAISS에 저장
+vectors = embed_model.encode(texts)
+dimension = vectors.shape[1]  # 임베딩 차원
+index = faiss.IndexFlatL2(dimension)  # L2 거리 기반 FAISS 인덱스 생성
+index.add(np.array(vectors))  # FAISS에 벡터 추가
+
+# 5. 검색 API (GPTs에서 호출 가능)
+@app.get("/search")
+def search(query: str):
+    """ 사용자의 쿼리를 벡터 변환 후, FAISS에서 검색하여 관련 문서 반환 """
+    query_vector = embed_model.encode([query])
+    query_vector = np.array(query_vector, dtype=np.float32)  # FAISS 호환
+
+    _, I = index.search(query_vector, k=3)  # FAISS로 Top-3 검색
+    results = [texts[i] for i in I[0]]  # 검색된 문서 반환
+
+    return {"retrieved_docs": results}