Update app.py

app.py

@@ -6,7 +6,6 @@ import os
 from threading import Thread
 import random
 from datasets import load_dataset
-from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 
@@ -18,24 +17,45 @@ MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
 MODELS = os.environ.get("MODELS")
 MODEL_NAME = MODEL_ID.split("/")[-1]
 
-# 임베딩 모델 로드
-embedding_model = SentenceTransformer('sentence-transformers/xlm-r-100langs-bert-base-nli-stsb-mean-tokens')
+# 모델과 토크나이저 로드
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+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)
+    # 마지막 히든 스테이트의 평균을 임베딩으로 사용
+    embeddings = outputs.last_hidden_state.mean(dim=1)
+    return embeddings
+
 # 데이터셋의 질문들을 임베딩
 questions = wiki_dataset['train']['question'][:10000]  # 처음 10000개만 사용
-question_embeddings = embedding_model.encode(questions, convert_to_tensor=True)
+question_embeddings = []
+batch_size = 32
+
+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 = torch.cat(question_embeddings, dim=0)
 
 def find_relevant_context(query, top_k=3):
     # 쿼리 임베딩
-    query_embedding = embedding_model.encode(query, convert_to_tensor=True)
+    query_embedding = get_embeddings(query, model, tokenizer)
     
     # 코사인 유사도 계산
     similarities = cosine_similarity(
-        query_embedding.cpu().numpy().reshape(1, -1),
+        query_embedding.cpu().numpy(),
         question_embeddings.cpu().numpy()
     )[0]
     
@@ -47,7 +67,8 @@ def find_relevant_context(query, top_k=3):
     for idx in top_indices:
         relevant_contexts.append({
             'question': questions[idx],
-            'answer': wiki_dataset['train']['answer'][idx]
+            'answer': wiki_dataset['train']['answer'][idx],
+            'similarity': similarities[idx]
         })
     
     return relevant_contexts