Update app.py

app.py

@@ -33,21 +33,26 @@ def get_embeddings(text, model, tokenizer):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
     with torch.no_grad():
         outputs = model(**inputs)
-    # 마지막 히든 스테이트의 평균을 임베딩으로 사용
-    embeddings = outputs.last_hidden_state.mean(dim=1)
+    # hidden states의 평균을 임베딩으로 사용
+    hidden_states = outputs[0]  # 모델의 마지막 레이어 출력
+    embeddings = hidden_states.mean(dim=1)
     return embeddings
 
 # 데이터셋의 질문들을 임베딩
-questions = wiki_dataset['train']['question'][:10000]  # 처음 10000개만 사용
+print("임베딩 생성 시작...")
+questions = wiki_dataset['train']['question'][:1000]  # 처음 1000개만 사용 (테스트용)
 question_embeddings = []
-batch_size = 32
+batch_size = 8  # 배치 사이즈 줄임
 
 for i in range(0, len(questions), batch_size):
     batch = questions[i:i+batch_size]
     batch_embeddings = get_embeddings(batch, model, tokenizer)
-    question_embeddings.append(batch_embeddings)
+    question_embeddings.append(batch_embeddings.cpu())
+    if i % 100 == 0:
+        print(f"Processed {i}/{len(questions)} questions")
 
 question_embeddings = torch.cat(question_embeddings, dim=0)
+print("임베딩 생성 완료")
 
 def find_relevant_context(query, top_k=3):
     # 쿼리 임베딩
@@ -56,7 +61,7 @@ def find_relevant_context(query, top_k=3):
     # 코사인 유사도 계산
     similarities = cosine_similarity(
         query_embedding.cpu().numpy(),
-        question_embeddings.cpu().numpy()
+        question_embeddings.numpy()
     )[0]
     
     # 가장 유사한 질문들의 인덱스