Update app.py

app.py

@@ -1,403 +1,166 @@
-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 gradio as gr
 import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import os
 from threading import Thread
+import random
+from datasets import load_dataset
 
-# 환경 변수 설정
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 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.float16,
-                device_map="auto",
-                trust_remote_code=True,
-                low_cpu_mem_usage=True
-            )
-            print("모델 로딩 완료")
-            
-        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: "
-
-            # 토크나이징 및 device 설정
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=4096
-            )
-            
-            # 모든 텐서를 GPU로 이동
-            input_ids = inputs.input_ids.to(self.model.device)
-            attention_mask = inputs.attention_mask.to(self.model.device)
-
-            # 생성
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    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
-                )
-
-            # 디코딩 전에 CPU로 이동
-            outputs = outputs.cpu()
-            generated_text = self.tokenizer.decode(
-                outputs[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 "❌ 데이터셋 구조 분석 실패"
+MODELS = os.environ.get("MODELS")
+MODEL_NAME = MODEL_ID.split("/")[-1]
+
+TITLE = "<h1><center>새로운 일본어 LLM 모델 웹 UI</center></h1>"
+
+DESCRIPTION = f"""
+<h3>모델: <a href="https://huggingface.co/CohereForAI/c4ai-command-r7b-12-2024">CohereForAI/c4ai-command-r7b-12-2024</a></h3>
+<center>
+<p>
+<br>
+cc-by-nc
+</p>
+</center>
+"""
+
+CSS = """
+.duplicate-button {
+    margin: auto !important;
+    color: white !important;
+    background: black !important;
+    border-radius: 100vh !important;
+}
+h3 {
+    text-align: center;
+}
+.chatbox .messages .message.user {
+    background-color: #e1f5fe;
+}
+.chatbox .messages .message.bot {
+    background-color: #eeeeee;
+}
+"""
+
+# 모델과 토크나이저 로드
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# 데이터셋 로드
+dataset = load_dataset("elyza/ELYZA-tasks-100")
+print(dataset)
+
+split_name = "train" if "train" in dataset else "test"
+examples_list = list(dataset[split_name])
+examples = random.sample(examples_list, 50)
+example_inputs = [[example['input']] for example in examples]
+
+@spaces.GPU
+def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
+    print(f'message is - {message}')
+    print(f'history is - {history}')
+    conversation = []
+    for prompt, answer in history:
+        conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(input_ids, return_tensors="pt").to(0)
     
-    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})"
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+
+    generate_kwargs = dict(
+        inputs, 
+        streamer=streamer,
+        top_k=top_k,
+        top_p=top_p,
+        repetition_penalty=penalty,
+        max_new_tokens=max_new_tokens, 
+        do_sample=True, 
+        temperature=temperature,
+        eos_token_id=[255001],
+    )
     
-    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]]
-
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", title="GiniGEN 🤖") as demo:
-    initial_history = chat_history.format_for_display()
-    with gr.Row():
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(
-                value=initial_history,
-                height=600, 
-                label="대화창 💬",
-                show_label=True
-            )    
-
-            msg = gr.Textbox(
-                label="메시지 입력",
-                show_label=False,
-                placeholder="무엇이든 물어보세요... 💭",
-                container=False
-            )
-            with gr.Row():
-                clear = gr.ClearButton([msg, chatbot], value="대화내용 지우기")
-                send = gr.Button("보내기 📤")
-        
-        with gr.Column(scale=1):
-            gr.Markdown("### GiniGEN 🤖 [파일 업로드] 📁\n지원 형식: 텍스트, 코드, CSV, Parquet 파일")
-            file_upload = gr.File(
-                label="파일 선택",
-                file_types=["text", ".csv", ".parquet"],
-                type="filepath"
-            )
-            
-            with gr.Accordion("고급 설정 ⚙️", open=False):
-                system_message = gr.Textbox(label="시스템 메시지 📝", value="")
-                max_tokens = gr.Slider(minimum=1, maximum=8000, value=4000, label="최대 토큰 수 📊")
-                temperature = gr.Slider(minimum=0, maximum=1, value=0.7, label="창의성 수준 🌡️")
-                top_p = gr.Slider(minimum=0, maximum=1, value=0.9, label="응답 다양성 📈")
-
-    gr.Examples(
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
+
+chatbot = gr.Chatbot(height=500)
+
+with gr.Blocks(css=CSS) as demo:
+    gr.HTML(TITLE)
+    gr.HTML(DESCRIPTION)
+    gr.ChatInterface(
+        fn=stream_chat,
+        chatbot=chatbot,
+        fill_height=True,
+        theme="soft",
+        additional_inputs_accordion=gr.Accordion(label="⚙️ 매개변수", open=False, render=False),
+        additional_inputs=[
+            gr.Slider(
+                minimum=0,
+                maximum=1,
+                step=0.1,
+                value=0.8,
+                label="온도",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=128,
+                maximum=1000000,
+                step=1,
+                value=100000,
+                label="최대 토큰 수",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.1,
+                value=0.8,
+                label="상위 확률",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=20,
+                step=1,
+                value=20,
+                label="상위 K",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=2.0,
+                step=0.1,
+                value=1.0,
+                label="반복 패널티",
+                render=False,
+            ),
+        ],
         examples=[
-            ["안녕하세요! 어떤 도움이 필요하신가요? 🤝"],
-            ["제가 이해하기 쉽게 설명해 주시겠어요? 📚"],
-            ["이 내용을 실제로 어떻게 활용할 수 있을까요? 🎯"],
-            ["추가로 조언해 주실 내용이 있으신가요? ✨"],
-            ["궁금한 점이 더 있는데 여쭤봐도 될까요? 🤔"],
+            ["아이의 여름방학 과학 프로젝트를 위한 5가지 아이디어를 주세요."],
+            ["마크다운을 사용하여 브레이크아웃 게임 만들기 튜토리얼을 작성해주세요."],
+            ["초능력을 가진 주인공의 SF 이야기 시나리오를 작성해주세요. 복선 설정, 테마와 로그라인을 논리적으로 사용해주세요"],
+            ["아이의 여름방학 자유연구를 위한 5가지 아이디어와 그 방법을 간단히 알려주세요."],
+            ["퍼즐 게임 스크립트 작성을 위한 조언 부탁드립니다"],
+            ["마크다운 형식으로 블록 깨기 게임 제작 교과서를 작성해주세요"],
+            ["실버 川柳를 생각해주세요"],
+            ["일본어 관용구, 속담에 관한 시험 문제를 만들어주세요"],
+            ["도라에몽의 등장인물을 알려주세요"],
+            ["오코노미야키 만드는 방법을 알려주세요"],
+            ["문제 9.11과 9.9 중 어느 것이 더 큰가요? step by step으로 논리적으로 생각해주세요."],
         ],
-        inputs=msg,
-    )
-
-    def clear_chat():
-        chat_history.clear_history()
-        return None, None
-
-    msg.submit(
-        chat,
-        inputs=[msg, chatbot, file_upload, system_message, max_tokens, temperature, top_p],
-        outputs=[msg, chatbot]
-    )
-
-    send.click(
-        chat,
-        inputs=[msg, chatbot, file_upload, system_message, max_tokens, temperature, top_p],
-        outputs=[msg, chatbot]
-    )
-
-    clear.click(
-        clear_chat,
-        outputs=[msg, chatbot]
-    )
-
-    file_upload.change(
-        lambda: "파일 분석을 시작합니다...",
-        outputs=msg
-    ).then(
-        chat,
-        inputs=[msg, chatbot, file_upload, system_message, max_tokens, temperature, top_p],
-        outputs=[msg, chatbot]
+        cache_examples=False,
     )
 
 if __name__ == "__main__":
-    demo.launch(
-        share=False,
-        show_error=True,
-        server_name="0.0.0.0",
-        server_port=7860,
-        max_threads=1
-    )
+    demo.launch()