Update app.py

app.py

@@ -6,8 +6,8 @@ import os
 from threading import Thread
 import random
 from datasets import load_dataset
-from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
 
 # GPU 메모리 관리
 torch.cuda.empty_cache()
@@ -29,40 +29,19 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 wiki_dataset = load_dataset("lcw99/wikipedia-korean-20240501-1million-qna")
 print("Wikipedia dataset loaded:", wiki_dataset)
 
-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)
-    # hidden states의 평균을 임베딩으로 사용
-    hidden_states = outputs[0]  # 모델의 마지막 레이어 출력
-    embeddings = hidden_states.mean(dim=1)
-    return embeddings
-
-# 데이터셋의 질문들을 임베딩
-print("임베딩 생성 시작...")
-questions = wiki_dataset['train']['question'][:1000]  # 처음 1000개만 사용 (테스트용)
-question_embeddings = []
-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.cpu())
-    if i % 100 == 0:
-        print(f"Processed {i}/{len(questions)} questions")
-
-question_embeddings = torch.cat(question_embeddings, dim=0)
-print("임베딩 생성 완료")
+# TF-IDF 벡터라이저 초기화 및 학습
+print("TF-IDF 벡터화 시작...")
+questions = wiki_dataset['train']['question'][:10000]  # 처음 10000개만 사용
+vectorizer = TfidfVectorizer(max_features=1000)
+question_vectors = vectorizer.fit_transform(questions)
+print("TF-IDF 벡터화 완료")
 
 def find_relevant_context(query, top_k=3):
-    # 쿼리 임베딩
-    query_embedding = get_embeddings(query, model, tokenizer)
+    # 쿼리 벡터화
+    query_vector = vectorizer.transform([query])
     
     # 코사인 유사도 계산
-    similarities = cosine_similarity(
-        query_embedding.cpu().numpy(),
-        question_embeddings.numpy()
-    )[0]
+    similarities = (query_vector * question_vectors.T).toarray()[0]
     
     # 가장 유사한 질문들의 인덱스
     top_indices = np.argsort(similarities)[-top_k:][::-1]
@@ -70,11 +49,12 @@ def find_relevant_context(query, top_k=3):
     # 관련 컨텍스트 추출
     relevant_contexts = []
     for idx in top_indices:
-        relevant_contexts.append({
-            'question': questions[idx],
-            'answer': wiki_dataset['train']['answer'][idx],
-            'similarity': similarities[idx]
-        })
+        if similarities[idx] > 0:  # 유사도가 0보다 큰 경우만 포함
+            relevant_contexts.append({
+                'question': questions[idx],
+                'answer': wiki_dataset['train']['answer'][idx],
+                'similarity': similarities[idx]
+            })
     
     return relevant_contexts
 
@@ -83,11 +63,11 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
     print(f'message is - {message}')
     print(f'history is - {history}')
     
-    # RAG: 관련 컨텍스트 찾기
+    # 관련 컨텍스트 찾기
     relevant_contexts = find_relevant_context(message)
     context_prompt = "\n\n관련 참고 정보:\n"
     for ctx in relevant_contexts:
-        context_prompt += f"Q: {ctx['question']}\nA: {ctx['answer']}\n\n"
+        context_prompt += f"Q: {ctx['question']}\nA: {ctx['answer']}\n유사도: {ctx['similarity']:.3f}\n\n"
     
     # 대화 히스토리 구성
     conversation = []
@@ -97,6 +77,7 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
             {"role": "assistant", "content": answer}
         ])
     
+
     # 컨텍스트를 포함한 최종 프롬프트 구성
     final_message = context_prompt + "\n현재 질문: " + message
     conversation.append({"role": "user", "content": final_message})