Update app.py

app.py

@@ -6,62 +6,79 @@ 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
+
+# GPU 메모리 관리
+torch.cuda.empty_cache()
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
 MODELS = os.environ.get("MODELS")
 MODEL_NAME = MODEL_ID.split("/")[-1]
 
-TITLE = "<h1><center>온디바이스 AI(Open LLM 모델)</center></h1>"
-
-CSS = """
-.duplicate-button {
-    margin: auto !important;
-    color: white !important;
-    background: black !important;
-    border-radius: 100vh !important;
-}
-h3 {
-    text-align: center;
-}
-.chatbox .messages .message.user {
-    background-color: #e1f5fe;
-}
-.chatbox .messages .message.bot {
-    background-color: #eeeeee;
-}
-"""
+# 임베딩 모델 로드
+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)
 
-# 데이터셋 로드
-dataset = load_dataset("elyza/ELYZA-tasks-100")
-print(dataset)
+# 데이터셋의 질문들을 임베딩
+questions = wiki_dataset['train']['question'][:10000]  # 처음 10000개만 사용
+question_embeddings = embedding_model.encode(questions, convert_to_tensor=True)
 
-split_name = "train" if "train" in dataset else "test"
-examples_list = list(dataset[split_name])
-examples = random.sample(examples_list, 50)
-example_inputs = [[example['input']] for example in examples]
+def find_relevant_context(query, top_k=3):
+    # 쿼리 임베딩
+    query_embedding = embedding_model.encode(query, convert_to_tensor=True)
+    
+    # 코사인 유사도 계산
+    similarities = cosine_similarity(
+        query_embedding.cpu().numpy().reshape(1, -1),
+        question_embeddings.cpu().numpy()
+    )[0]
+    
+    # 가장 유사한 질문들의 인덱스
+    top_indices = np.argsort(similarities)[-top_k:][::-1]
+    
+    # 관련 컨텍스트 추출
+    relevant_contexts = []
+    for idx in top_indices:
+        relevant_contexts.append({
+            'question': questions[idx],
+            'answer': wiki_dataset['train']['answer'][idx]
+        })
+    
+    return relevant_contexts
 
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
     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"
+    
+    # 대화 히스토리 구성
     conversation = []
     for prompt, answer in history:
-        conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
-    conversation.append({"role": "user", "content": message})
+        conversation.extend([
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": answer}
+        ])
+    
+    # 컨텍스트를 포함한 최종 프롬프트 구성
+    final_message = context_prompt + "\n현재 질문: " + message
+    conversation.append({"role": "user", "content": final_message})
 
     input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(input_ids, return_tensors="pt").to(0)
     
-    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, timeout
 
     generate_kwargs = dict(
         inputs,