Update app.py

app.py

@@ -3,17 +3,18 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 from langchain.vectorstores import Chroma
 import os
+import psutil
+import time
 
 # Hugging Face 모델 ID
 model_id = "hewoo/hehehehe"
-token = os.getenv("HF_API_TOKEN")  # 필요한 경우 사용자에게 Hugging Face API 토큰 입력을 요청할 수 있음
+# 모델과 토크나이저 로드 (토큰 없이 사용)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
-# 모델과 토크나이저 로드
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
-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)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.3, top_p=0.85, top_k=40, repetition_penalty=1.2)
 
 # 사용자 정의 임베딩 클래스 생성
 class CustomEmbedding:
@@ -27,29 +28,49 @@ class CustomEmbedding:
         return [self.model.encode(text, convert_to_tensor=True).tolist() for text in texts]
 
 # 임베딩 모델 및 벡터 스토어 설정
-embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+embedding_model = SentenceTransformer("jhgan/ko-sroberta-multitask")
 embedding_function = CustomEmbedding(embedding_model)
 
 # Chroma 벡터 스토어 설정
-persist_directory = "./chroma_batch_vectors"  # Spaces 환경에 맞게 조정 필요
+persist_directory = "./chroma_batch_vectors"
 vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embedding_function)
 retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
 
 # 질문에 대한 응답 생성 함수
 def generate_response(user_input):
+    start_time = time.time()  # 시작 시간 기록
+
+    # 문서 검색 및 맥락 생성
     search_results = retriever.get_relevant_documents(user_input)
     context = "\n".join([result.page_content for result in search_results])
-    input_text = f"맥락: {context}\n질문: {user_input}"
+    input_text = f"""아래는 한국어로만 답변하는 어시스턴트입니다.
+    사용자의 질문에 대해 제공된 맥락을 바탕으로 정확하고 자세한 답변을 한국어로 작성하세요.
+
+    맥락: {context}
+    질문: {user_input}
+    답변:"""
+
+    # 응답 생성
     response = pipe(input_text)[0]["generated_text"]
-    return response
+
+    end_time = time.time()  # 끝난 시간 기록
+    response_time = end_time - start_time  # 응답 시간 계산
+
+    # 메모리 사용량 모니터링
+    memory_info = psutil.virtual_memory()
+    memory_usage = memory_info.percent  # 메모리 사용량 비율(%)
+
+    return response, response_time, memory_usage
 
 # Streamlit 앱 UI
-st.title("챗봇 test")
+st.title("챗봇 데모")
 st.write("Llama 3.2-3B 모델을 사용한 챗봇입니다. 질문을 입력해 주세요.")
 
 # 사용자 입력 받기
 user_input = st.text_input("질문")
 if user_input:
-    response = generate_response(user_input)
+    response, response_time, memory_usage = generate_response(user_input)
     st.write("챗봇 응답:", response)
+    st.write(f"응답 시간: {response_time:.2f}초")
+    st.write(f"현재 메모리 사용량: {memory_usage}%")