Update app.py

app.py

@@ -2,19 +2,23 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 from langchain.vectorstores import Chroma
-import os
-import psutil
 import gc
+import psutil
 
-# Hugging Face 모델 ID
+# 모델 ID (공개된 모델이어야 함)
 model_id = "hewoo/hehehehe"
 
-# 캐시를 사용하여 모델과 파이프라인 로드
+# 메모리 모니터링 함수
+def monitor_memory():
+    memory_info = psutil.virtual_memory()
+    st.write(f"현재 메모리 사용량: {memory_info.percent}%")
+
+# 캐시를 사용하여 모델 및 파이프라인 로드
 @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)
+    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:
@@ -22,18 +26,18 @@ class CustomEmbedding:
         self.model = model
 
     def embed_query(self, text):
-        return self.model.encode(text, convert_to_tensor=True).tolist()
+        return self.model.encode(text, convert_to_tensor=False).tolist()
 
     def embed_documents(self, texts):
-        return [self.model.encode(text, convert_to_tensor=True).tolist() for text in texts]
+        return [self.model.encode(text, convert_to_tensor=False).tolist() for text in texts]
 
-# 임베딩 모델 및 벡터 스토어 설정 (캐시 적용)
+# 한국어 임베딩 모델 및 벡터 스토어 설정
 @st.cache_resource
 def load_embedding_model():
     return SentenceTransformer("jhgan/ko-sroberta-multitask")
 
 @st.cache_resource
-def load_vectorstore(_embedding_model):
+def load_vectorstore(_embedding_model):  # _embedding_model로 이름 변경
     embedding_function = CustomEmbedding(_embedding_model)
     return Chroma(persist_directory="./chroma_batch_vectors", embedding_function=embedding_function)
 
@@ -56,8 +60,7 @@ def generate_response(user_input):
 답변:"""
 
     response = pipe(prompt)[0]["generated_text"]
-    response = response.split("답변:")[-1].strip()
-    return response, context
+    return response
 
 # 모델 및 임베딩 모델 로드
 pipe = load_model()
@@ -66,20 +69,19 @@ vectorstore = load_vectorstore(embedding_model)
 
 # Streamlit 앱 UI
 st.title("챗봇 데모")
-st.write("한국어 모델을 사용한 챗봇입니다. 질문을 입력해 주세요.")
+st.write("Llama 3.2-3B 모델을 사용한 챗봇입니다. 질문을 입력해 주세요.")
+
+monitor_memory()  # 메모리 사용량 확인
 
 # 사용자 입력 받기
 user_input = st.text_input("질문")
 if user_input:
-    response, context = generate_response(user_input)
+    response = generate_response(user_input)
     st.write("챗봇 응답:", response)
-    st.write("사용된 맥락:", context)
-
-    # 메모리 상태 모니터링
-    memory_usage = psutil.virtual_memory().used / (1024 ** 3)
-    st.write(f"현재 메모리 사용량: {memory_usage:.2f} GB")
+    monitor_memory()  # 메모리 상태 업데이트
 
     # 메모리 해제
     del response
     gc.collect()
 
+