Update app.py

app.py

@@ -33,6 +33,40 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# 전역 변수
+model = None
+tokenizer = None
+current_file_context = None
+
+# CSS 스타일
+CSS = """
+.chat-container {
+    height: 600px !important;
+    margin-bottom: 10px;
+}
+.input-container {
+    height: 70px !important;
+    display: flex;
+    align-items: center;
+    gap: 10px;
+    margin-top: 5px;
+}
+.input-textbox {
+    height: 70px !important;
+    border-radius: 8px !important;
+    font-size: 1.1em !important;
+    padding: 10px 15px !important;
+}
+.custom-button {
+    background: linear-gradient(145deg, #2196f3, #1976d2);
+    color: white;
+    border-radius: 10px;
+    padding: 10px 20px;
+    font-weight: 600;
+    transition: all 0.3s ease;
+}
+"""
+
 # 설정 클래스
 class Config:
     def __init__(self):
@@ -55,8 +89,29 @@ class ChatResponse(BaseModel):
     status: str
     timestamp: datetime
 
+def initialize_model_and_tokenizer():
+    global model, tokenizer
+    try:
+        model = load_model()
+        tokenizer = AutoTokenizer.from_pretrained(config.MODEL_ID)
+        return True
+    except Exception as e:
+        logger.error(f"Initialization error: {str(e)}")
+        return False
+
 # 파일 처리 클래스
 class FileProcessor:
+    @staticmethod
+    def safe_file_read(file_path):
+        encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1']
+        for encoding in encodings:
+            try:
+                with open(file_path, 'r', encoding=encoding) as f:
+                    return f.read()
+            except UnicodeDecodeError:
+                continue
+        raise FileProcessingError("Unable to read file with supported encodings")
+
     @staticmethod
     def process_pdf(file_path):
         try:
@@ -78,22 +133,34 @@ class FileProcessor:
     @staticmethod
     def process_csv(file_path):
         try:
-            encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1']
-            for encoding in encodings:
-                try:
-                    return pd.read_csv(file_path, encoding=encoding)
-                except UnicodeDecodeError:
-                    continue
-            raise FileProcessingError("Unable to read CSV with supported encodings")
+            return pd.read_csv(file_path)
         except Exception as e:
             raise FileProcessingError(f"CSV processing error: {str(e)}")
 
+    @staticmethod
+    def analyze_file_content(content, file_type):
+        try:
+            if file_type == 'pdf':
+                words = len(content.split())
+                lines = content.count('\n') + 1
+                return f"PDF Analysis:\nWords: {words}\nLines: {lines}"
+            elif file_type == 'csv':
+                df = pd.DataFrame(content)
+                return f"CSV Analysis:\nRows: {len(df)}\nColumns: {len(df.columns)}"
+            else:
+                lines = content.split('\n')
+                return f"Text Analysis:\nLines: {len(lines)}"
+        except Exception as e:
+            raise FileProcessingError(f"Content analysis error: {str(e)}")
+
 # 메모리 관리
 @torch.no_grad()
 def clear_cuda_memory():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         gc.collect()
+        if model is not None:
+            model.cpu()
 
 # 모델 로드
 @spaces.GPU
@@ -129,32 +196,19 @@ def find_relevant_context(query, top_k=3):
     except Exception as e:
         logger.error(f"Context search error: {str(e)}")
         return []
-
 # 스트리밍 채팅
 @spaces.GPU
 def stream_chat(message: str, history: list, uploaded_file, temperature: float, 
                 max_new_tokens: int, top_p: float, top_k: int, penalty: float) -> Iterator[Tuple[str, list]]:
     """
     스트리밍 채팅 응답을 생성합니다.
-    
-    Args:
-        message (str): 사용자 입력 메시지
-        history (list): 대화 히스토리
-        uploaded_file: 업로드된 파일
-        temperature (float): 생성 온도
-        max_new_tokens (int): 최대 토큰 수
-        top_p (float): 상위 p 샘플링
-        top_k (int): 상위 k 샘플링
-        penalty (float): 반복 페널티
-        
-    Returns:
-        Iterator[Tuple[str, list]]: 생성된 응답과 업데이트된 히스토리
     """
-    global model, current_file_context
+    global model, tokenizer, current_file_context
     
     try:
-        if model is None:
-            model = load_model()
+        if model is None or tokenizer is None:
+            if not initialize_model_and_tokenizer():
+                raise Exception("Model initialization failed")
             
         logger.info(f'Processing message: {message}')
         logger.debug(f'History length: {len(history)}')
@@ -169,9 +223,9 @@ def stream_chat(message: str, history: list, uploaded_file, temperature: float,
                 elif file_ext == '.csv':
                     content = FileProcessor.process_csv(uploaded_file.name)
                 else:
-                    content = safe_file_read(uploaded_file.name)
+                    content = FileProcessor.safe_file_read(uploaded_file.name)
                 
-                file_context = analyze_file_content(content, file_ext)
+                file_context = FileProcessor.analyze_file_content(content, file_ext.replace('.', ''))
                 current_file_context = file_context
             except Exception as e:
                 logger.error(f"File processing error: {str(e)}")
@@ -199,7 +253,16 @@ def stream_chat(message: str, history: list, uploaded_file, temperature: float,
             return_tensors="pt"
         ).to("cuda")
 
-        streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+        # 입력 길이 체크
+        if len(inputs.input_ids[0]) > config.MAX_TOKENS:
+            raise ValueError("Input too long")
+
+        streamer = TextIteratorStreamer(
+            tokenizer, 
+            timeout=30.0,
+            skip_prompt=True, 
+            skip_special_tokens=True
+        )
         
         generate_kwargs = dict(
             inputs, 
@@ -215,13 +278,14 @@ def stream_chat(message: str, history: list, uploaded_file, temperature: float,
         
         clear_cuda_memory()
         
-        thread = Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
+        with torch.no_grad():
+            thread = Thread(target=model.generate, kwargs=generate_kwargs)
+            thread.start()
 
-        buffer = ""
-        for new_text in streamer:
-            buffer += new_text
-            yield "", history + [[message, buffer]]
+            buffer = ""
+            for new_text in streamer:
+                buffer += new_text
+                yield "", history + [[message, buffer]]
 
         clear_cuda_memory()
 
@@ -232,8 +296,7 @@ def stream_chat(message: str, history: list, uploaded_file, temperature: float,
 
 # UI 생성
 def create_demo():
-    with gr.Blocks(css=UPDATED_CSS) as demo:
-        # UI 컴포넌트 구성
+    with gr.Blocks(css=CSS) as demo:
         with gr.Column(elem_classes="markdown-style"):
             gr.Markdown("""
                 # 🤖 RAGOndevice
@@ -244,11 +307,10 @@ def create_demo():
         chatbot = gr.Chatbot(
             value=[],
             height=600,
-            label="GiniGEN AI Assistant",
+            label="AI Assistant",
             elem_classes="chat-container"
         )
         
-        # 입력 컴포넌트
         with gr.Row(elem_classes="input-container"):
             with gr.Column(scale=1, min_width=70):
                 file_upload = gr.File(
@@ -283,7 +345,6 @@ def create_demo():
                     scale=1
                 )
         
-        # 고급 설정
         with gr.Accordion("🎮 Advanced Settings", open=False):
             with gr.Row():
                 with gr.Column(scale=1):
@@ -318,26 +379,43 @@ def create_demo():
 
 # 메인 실행
 if __name__ == "__main__":
-    # 위키피디아 데이터셋 로드
-    wiki_dataset = load_dataset("lcw99/wikipedia-korean-20240501-1million-qna")
-    logger.info("Wikipedia dataset loaded")
+    try:
+        # 모델 초기화
+        if not initialize_model_and_tokenizer():
+            logger.error("Failed to initialize model and tokenizer")
+            exit(1)
+            
+        # 위키피디아 데이터셋 로드
+        wiki_dataset = load_dataset("lcw99/wikipedia-korean-20240501-1million-qna")
+        logger.info("Wikipedia dataset loaded")
 
-    # TF-IDF 벡터라이저 초기화
-    questions = wiki_dataset['train']['question'][:10000]
-    vectorizer = TfidfVectorizer(max_features=1000)
-    question_vectors = vectorizer.fit_transform(questions)
-    logger.info("TF-IDF vectorization completed")
+        # TF-IDF 벡터라이저 초기화
+        questions = wiki_dataset['train']['question'][:10000]
+        vectorizer = TfidfVectorizer(max_features=1000)
+        question_vectors = vectorizer.fit_transform(questions)
+        logger.info("TF-IDF vectorization completed")
 
-    # UI 실행
-    demo = create_demo()
-    demo.launch()
+        # UI 실행
+        demo = create_demo()
+        demo.launch(share=False, server_name="0.0.0.0")
+        
+    except Exception as e:
+        logger.error(f"Application startup error: {str(e)}")
+        exit(1)
 
 # 테스트 코드
 class TestChatBot(unittest.TestCase):
+    def setUp(self):
+        self.file_processor = FileProcessor()
+
     def test_file_processing(self):
-        # 테스트 구현
-        pass
+        # 파일 처리 테스트
+        test_content = "Test content"
+        result = self.file_processor.analyze_file_content(test_content, 'txt')
+        self.assertIsNotNone(result)
 
     def test_context_search(self):
-        # 테스트 구현
-        pass
+        # 컨텍스트 검색 테스트
+        test_query = "테스트 질문"
+        result = find_relevant_context(test_query)
+        self.assertIsInstance(result, list)