Update app.py

app.py

@@ -2,68 +2,73 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 from langchain.vectorstores import Chroma
-import os
+import gc
+import psutil
 
-# Hugging Face 모델 ID
-model_id = "hewoo/hehehehe"  # 업로드한 모델의 repo_id
-token = os.getenv("HF_API_TOKEN")  # Hugging Face API 토큰 (필요 시 설정)
+# 모델 ID (공개된 모델이어야 함)
+model_id = "hewoo/hehehehe"
 
-# 모델과 토크나이저 로드
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
-model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+# 메모리 모니터링 함수
+def monitor_memory():
+    memory_info = psutil.virtual_memory()
+    st.write(f"현재 메모리 사용량: {memory_info.percent}%")
 
-# 텍스트 생성 파이프라인 설정
-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)
+# 캐시를 사용하여 모델 및 파이프라인 로드
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+    return 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()
+        return self.model.encode(text, convert_to_tensor=True).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)
+        return [self.model.encode(text, convert_to_tensor=True).tolist() for text in texts]
 
-# Chroma 벡터 스토어 설정
-persist_directory = "./chroma_batch_vectors"
-vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embedding_function)
+# 임베딩 모델 및 벡터 스토어 설정
+@st.cache_resource
+def load_embedding_model():
+    return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# 검색 기능 설정
-retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+@st.cache_resource
+def load_vectorstore(embedding_model):
+    embedding_function = CustomEmbedding(embedding_model)
+    return Chroma(persist_directory="./chroma_batch_vectors", embedding_function=embedding_function)
 
-
-# 검색 결과 요약 함수
-def summarize_results(search_results):
-    combined_text = "\n".join([result.page_content for result in search_results])
-    summary = summarizer(combined_text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
-    return summary
-
-# 검색 및 응답 생성 함수
+# 질문에 대한 응답 생성 함수
 def generate_response(user_input):
-    # 검색 및 맥락 생성
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
     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}"
     response = pipe(input_text)[0]["generated_text"]
-
     return response
 
+# 모델 및 임베딩 모델 로드
+pipe = load_model()
+embedding_model = load_embedding_model()
+vectorstore = load_vectorstore(embedding_model)
+
 # Streamlit 앱 UI
-st.title("챗봇테스트")
+st.title("챗봇 데모")
 st.write("Llama 3.2-3B 모델을 사용한 챗봇입니다. 질문을 입력해 주세요.")
 
+monitor_memory()  # 메모리 사용량 확인
+
 # 사용자 입력 받기
 user_input = st.text_input("질문")
 if user_input:
     response = generate_response(user_input)
     st.write("챗봇 응답:", response)
+    monitor_memory()  # 메모리 상태 업데이트
+
+    # 메모리 해제
+    del response
+    gc.collect()
+