Update app.py

app.py

@@ -15,12 +15,31 @@ model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
 # 텍스트 생성 파이프라인 설정
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.5, top_p=0.85, top_k=40, repetition_penalty=1.2)
 
-# 임베딩 모델 및 검색 기능 설정
+
+
+# 사용자 정의 임베딩 클래스 생성
+class CustomEmbedding:
+    def __init__(self, model):
+        self.model = model
+
+    def embed_query(self, text):
+        return self.model.encode(text, convert_to_tensor=False).tolist()
+
+    def embed_documents(self, texts):
+        return [self.model.encode(text, convert_to_tensor=False).tolist() for text in texts]
+
+# 임베딩 모델 설정
 embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+embedding_function = CustomEmbedding(embedding_model)
+
+# Chroma 벡터 스토어 설정
 persist_directory = "./chroma_batch_vectors"
-vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embedding_model.encode)
+vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embedding_function)
+
+# 검색 기능 설정
 retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
 
+
 # 검색 결과 요약 함수
 def summarize_results(search_results):
     combined_text = "\n".join([result.page_content for result in search_results])