Update app.py

app.py

@@ -5,7 +5,7 @@ import pandas as pd
 import json
 from datetime import datetime
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 from threading import Thread
 
@@ -17,8 +17,12 @@ class ModelManager:
     def __init__(self):
         self.tokenizer = None
         self.model = None
-        self.setup_model()
+        # 초기화는 첫 요청 시에 수행
     
+    def ensure_model_loaded(self):
+        if self.model is None or self.tokenizer is None:
+            self.setup_model()
+
     @spaces.GPU
     def setup_model(self):
         try:
@@ -42,13 +46,320 @@ class ModelManager:
                 trust_remote_code=True,
                 low_cpu_mem_usage=True
             )
-            self.model.eval()  # 평가 모드로 설정
+            self.model.eval()
             print("모델 로딩 완료")
             
-            # 모델과 토크나이저가 제대로 로드되었는지 확인
-            if self.model is None or self.tokenizer is None:
-                raise Exception("모델 또는 토크나이저가 제대로 초기화되지 않았습니다.")
+        except Exception as e:
+            print(f"모델 로딩 중 오류 발생: {e}")
+            raise Exception(f"모델 로딩 실패: {e}")
+
+    @spaces.GPU
+    def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
+        try:
+            # 모델이 로드되어 있는지 확인
+            self.ensure_model_loaded()
+
+            # 입력 텍스트 준비
+            prompt = ""
+            for msg in messages:
+                role = msg["role"]
+                content = msg["content"]
+                if role == "system":
+                    prompt += f"System: {content}\n"
+                elif role == "user":
+                    prompt += f"Human: {content}\n"
+                elif role == "assistant":
+                    prompt += f"Assistant: {content}\n"
+            prompt += "Assistant: "
+
+            # 입력 인코딩
+            input_ids = self.tokenizer.encode(
+                prompt,
+                return_tensors="pt",
+                add_special_tokens=True
+            ).to(self.model.device)
+
+            # 응답 생성
+            with torch.no_grad():
+                output_ids = self.model.generate(
+                    input_ids,
+                    max_new_tokens=max_tokens,
+                    do_sample=True,
+                    temperature=temperature,
+                    top_p=top_p,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    num_return_sequences=1
+                )
+
+            # 응답 디코딩
+            generated_text = self.tokenizer.decode(
+                output_ids[0][input_ids.shape[1]:],
+                skip_special_tokens=True
+            )
+
+            # 단어 단위로 스트리밍
+            words = generated_text.split()
+            for word in words:
+                yield type('Response', (), {
+                    'choices': [type('Choice', (), {
+                        'delta': {'content': word + " "}
+                    })()]
+                })()
+
+        except Exception as e:
+            print(f"응답 생성 중 오류 발생: {e}")
+            raise Exception(f"응답 생성 실패: {e}")
+
+class ChatHistory:
+    def __init__(self):
+        self.history = []
+        self.history_file = "/tmp/chat_history.json"
+        self.load_history()
+
+    def add_conversation(self, user_msg: str, assistant_msg: str):
+        conversation = {
+            "timestamp": datetime.now().isoformat(),
+            "messages": [
+                {"role": "user", "content": user_msg},
+                {"role": "assistant", "content": assistant_msg}
+            ]
+        }
+        self.history.append(conversation)
+        self.save_history()
+
+    def format_for_display(self):
+        formatted = []
+        for conv in self.history:
+            formatted.append([
+                conv["messages"][0]["content"],
+                conv["messages"][1]["content"]
+            ])
+        return formatted
+
+    def get_messages_for_api(self):
+        messages = []
+        for conv in self.history:
+            messages.extend([
+                {"role": "user", "content": conv["messages"][0]["content"]},
+                {"role": "assistant", "content": conv["messages"][1]["content"]}
+            ])
+        return messages
+
+    def clear_history(self):
+        self.history = []
+        self.save_history()
+
+    def save_history(self):
+        try:
+            with open(self.history_file, 'w', encoding='utf-8') as f:
+                json.dump(self.history, f, ensure_ascii=False, indent=2)
+        except Exception as e:
+            print(f"히스토리 저장 실패: {e}")
+
+    def load_history(self):
+        try:
+            if os.path.exists(self.history_file):
+                with open(self.history_file, 'r', encoding='utf-8') as f:
+                    self.history = json.load(f)
+        except Exception as e:
+            print(f"히스토리 로드 실패: {e}")
+            self.history = []
+
+# 전역 인스턴스 생성
+chat_history = ChatHistory()
+model_manager = ModelManager()
+
+def analyze_file_content(content, file_type):
+    """Analyze file content and return structural summary"""
+    if file_type in ['parquet', 'csv']:
+        try:
+            lines = content.split('\n')
+            header = lines[0]
+            columns = header.count('|') - 1
+            rows = len(lines) - 3
+            return f"📊 데이터셋 구조: {columns}개 컬럼, {rows}개 데이터"
+        except:
+            return "❌ 데이터셋 구조 분석 실패"
+    
+    lines = content.split('\n')
+    total_lines = len(lines)
+    non_empty_lines = len([line for line in lines if line.strip()])
+    
+    if any(keyword in content.lower() for keyword in ['def ', 'class ', 'import ', 'function']):
+        functions = len([line for line in lines if 'def ' in line])
+        classes = len([line for line in lines if 'class ' in line])
+        imports = len([line for line in lines if 'import ' in line or 'from ' in line])
+        return f"💻 코드 구조: {total_lines}줄 (함수: {functions}, 클래스: {classes}, 임포트: {imports})"
+    
+    paragraphs = content.count('\n\n') + 1
+    words = len(content.split())
+    return f"📝 문서 구조: {total_lines}줄, {paragraphs}단락, 약 {words}단어"
+
+def read_uploaded_file(file):
+    if file is None:
+        return "", ""
+    try:
+        file_ext = os.path.splitext(file.name)[1].lower()
+        
+        if file_ext == '.parquet':
+            df = pd.read_parquet(file.name, engine='pyarrow')
+            content = df.head(10).to_markdown(index=False)
+            return content, "parquet"
+        elif file_ext == '.csv':
+            encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1']
+            for encoding in encodings:
+                try:
+                    df = pd.read_csv(file.name, encoding=encoding)
+                    content = f"📊 데이터 미리보기:\n{df.head(10).to_markdown(index=False)}\n\n"
+                    content += f"\n📈 데이터 정보:\n"
+                    content += f"- 전체 행 수: {len(df)}\n"
+                    content += f"- 전체 열 수: {len(df.columns)}\n"
+                    content += f"- 컬럼 목록: {', '.join(df.columns)}\n"
+                    content += f"\n📋 컬럼 데이터 타입:\n"
+                    for col, dtype in df.dtypes.items():
+                        content += f"- {col}: {dtype}\n"
+                    null_counts = df.isnull().sum()
+                    if null_counts.any():
+                        content += f"\n⚠️ 결측치:\n"
+                        for col, null_count in null_counts[null_counts > 0].items():
+                            content += f"- {col}: {null_count}개 누락\n"
+                    return content, "csv"
+                except UnicodeDecodeError:
+                    continue
+            raise UnicodeDecodeError(f"❌ 지원되는 인코딩으로 파일을 읽을 수 없습니다 ({', '.join(encodings)})")
+        else:
+            encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1']
+            for encoding in encodings:
+                try:
+                    with open(file.name, 'r', encoding=encoding) as f:
+                        content = f.read()
+                    return content, "text"
+                except UnicodeDecodeError:
+                    continue
+            raise UnicodeDecodeError(f"❌ 지원되는 인코딩으로 파일을 읽을 수 없습니다 ({', '.join(encodings)})")
+    except Exception as e:
+        return f"❌ 파일 읽기 오류: {str(e)}", "error"
+
+def chat(message, history, uploaded_file, system_message="", max_tokens=4000, temperature=0.7, top_p=0.9):
+    if not message:
+        return "", history
+
+    system_prefix = """저는 여러분의 친근하고 지적인 AI 어시스턴트 'GiniGEN'입니다.. 다음과 같은 원칙으로 소통하겠습니다:
+1. 🤝 친근하고 공감적인 태도로 대화
+2. 💡 명확하고 이해하기 쉬운 설명 제공
+3. 🎯 질문의 의도를 정확히 파악하여 맞춤형 답변
+4. 📚 필요한 경우 업로드된 파일 내용을 참고하여 구체적인 도움 제공
+5. ✨ 추가적인 통찰과 제안을 통한 가치 있는 대화
+항상 예의 바르고 친절하게 응답하며, 필요한 경우 구체적인 예시나 설명을 추가하여 
+이해를 돕겠습니다."""
+
+    try:
+        # 첫 메시지일 때 모델 로딩
+        model_manager.ensure_model_loaded()
+
+        if uploaded_file:
+            content, file_type = read_uploaded_file(uploaded_file)
+            if file_type == "error":
+                error_message = content
+                chat_history.add_conversation(message, error_message)
+                return "", history + [[message, error_message]]
+            
+            file_summary = analyze_file_content(content, file_type)
+            
+            if file_type in ['parquet', 'csv']:
+                system_message += f"\n\n파일 내용:\n```markdown\n{content}\n```"
+            else:
+                system_message += f"\n\n파일 내용:\n```\n{content}\n```"
                 
+            if message == "파일 분석을 시작합니다...":
+                message = f"""[파일 구조 분석] {file_summary}
+다음 관점에서 도움을 드리겠습니다:
+1. 📋 전반적인 내용 파악
+2. 💡 주요 특징 설명
+3. 🎯 실용적인 활용 방안
+4. ✨ 개선 제안
+5. 💬 추가 질문이나 필요한 설명"""
+
+        messages = [{"role": "system", "content": system_prefix + system_message}]
+        
+        if history:
+            for user_msg, assistant_msg in history:
+                messages.append({"role": "user", "content": user_msg})
+                messages.append({"role": "assistant", "content": assistant_msg})
+        
+        messages.append({"role": "user", "content": message})
+
+        partial_message = ""
+        
+        for response in model_manager.generate_response(
+            messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p
+        ):
+            token = response.choices[0].delta.get('content', '')
+            if token:
+                partial_message += token
+                current_history = history + [[message, partial_message]]
+                yield "", current_history
+
+        chat_history.add_conversation(message, partial_message)
+        
+    except Exception as e:
+        error_msg = f"❌ 오류가 발생했습니다: {str(e)}"
+        chat_history.add_conversation(message, error_msg)
+        yield "", history + [[message, error_msg]]import os
+from dotenv import load_dotenv
+import gradio as gr
+import pandas as pd
+import json
+from datetime import datetime
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import spaces
+from threading import Thread
+
+# 환경 변수 설정
+HF_TOKEN = os.getenv("HF_TOKEN")
+MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
+
+class ModelManager:
+    def __init__(self):
+        self.tokenizer = None
+        self.model = None
+        # 초기화는 첫 요청 시에 수행
+    
+    def ensure_model_loaded(self):
+        if self.model is None or self.tokenizer is None:
+            self.setup_model()
+
+    @spaces.GPU
+    def setup_model(self):
+        try:
+            print("토크나이저 로딩 시작...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                use_fast=True,
+                token=HF_TOKEN,
+                trust_remote_code=True
+            )
+            if not self.tokenizer.pad_token:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            print("토크나이저 로딩 완료")
+            
+            print("모델 로딩 시작...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+            self.model.eval()
+            print("모델 로딩 완료")
+            
         except Exception as e:
             print(f"모델 로딩 중 오류 발생: {e}")
             raise Exception(f"모델 로딩 실패: {e}")
@@ -56,8 +367,8 @@ class ModelManager:
     @spaces.GPU
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
-            if self.model is None or self.tokenizer is None:
-                raise Exception("모델이 초기화되지 않았습니다.")
+            # 모델이 로드되어 있는지 확인
+            self.ensure_model_loaded()
 
             # 입력 텍스트 준비
             prompt = ""
@@ -70,7 +381,7 @@ class ModelManager:
                     prompt += f"Human: {content}\n"
                 elif role == "assistant":
                     prompt += f"Assistant: {content}\n"
-            prompt += "Assistant: "  # 응답 시작 프롬프트
+            prompt += "Assistant: "
 
             # 입력 인코딩
             input_ids = self.tokenizer.encode(
@@ -92,7 +403,7 @@ class ModelManager:
                     num_return_sequences=1
                 )
 
-            # 응답 디코딩 및 스트리밍
+            # 응답 디코딩
             generated_text = self.tokenizer.decode(
                 output_ids[0][input_ids.shape[1]:],
                 skip_special_tokens=True
@@ -255,6 +566,9 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
 이해를 돕겠습니다."""
 
     try:
+        # 첫 메시지일 때 모델 로딩
+        model_manager.ensure_model_loaded()
+
         if uploaded_file:
             content, file_type = read_uploaded_file(uploaded_file)
             if file_type == "error":