Update app.py

app.py

@@ -6,11 +6,6 @@ import os
 from threading import Thread
 import random
 from datasets import load_dataset
-import gc
-
-# GPU 메모리 관리
-torch.cuda.empty_cache()
-gc.collect()
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
@@ -37,97 +32,58 @@ h3 {
 }
 """
 
-# 디바이스 설정
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# 모델과 토크나이저 로드 with 에러 처리
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        low_cpu_mem_usage=True,
-    )
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-except Exception as e:
-    print(f"모델 로딩 중 오류 발생: {str(e)}")
-    raise
+# 모델과 토크나이저 로드
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# 데이터셋 로드 with 에러 처리
-try:
-    dataset = load_dataset("elyza/ELYZA-tasks-100")
-    print(dataset)
+# 데이터셋 로드
+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]
-except Exception as e:
-    print(f"데이터셋 로딩 중 오류 발생: {str(e)}")
-    examples = []
-    example_inputs = []
+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]
 
-def error_handler(func):
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except Exception as e:
-            print(f"Error in {func.__name__}: {str(e)}")
-            return "죄송합니다. 오류가 발생했습니다. 잠시 후 다시 시도해주세요."
-    return wrapper
-
-@error_handler
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
-    try:
-        print(f'message is - {message}')
-        print(f'history is - {history}')
-        
-        # GPU 메모리 정리
-        torch.cuda.empty_cache()
-        
-        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(device)
-
-        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],
-        )
-        
-        thread = Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
+    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)
+    
+    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],
+    )
+    
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
 
-        buffer = ""
-        for new_text in streamer:
-            buffer += new_text
-            yield buffer
-            
-    except Exception as e:
-        print(f"Stream chat error: {str(e)}")
-        yield "죄송합니다. 응답 생성 중 오류가 발생했습니다."
-    finally:
-        # 메모리 정리
-        torch.cuda.empty_cache()
-        gc.collect()
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
 
-chatbot = gr.Chatbot(height=500)                                        
+chatbot = gr.Chatbot(height=500)
 
 CSS = """
 /* 전체 페이지 스타일링 */
@@ -136,7 +92,6 @@ body {
     min-height: 100vh;
     font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
 }
-
 /* 메인 컨테이너 */
 .container {
     max-width: 1200px;
@@ -149,7 +104,6 @@ body {
     transform: perspective(1000px) translateZ(0);
     transition: all 0.3s ease;
 }
-
 /* 제목 스타일링 */
 h1 {
     color: #2d3436;
@@ -159,14 +113,12 @@ h1 {
     text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
     transform: perspective(1000px) translateZ(20px);
 }
-
 h3 {
     text-align: center;
     color: #2d3436;
     font-size: 1.5rem;
     margin: 1rem 0;
 }
-
 /* 채팅박스 스타일링 */
 .chatbox {
     background: white;
@@ -179,7 +131,6 @@ h3 {
     transform: translateZ(0);
     transition: all 0.3s ease;
 }
-
 /* 메시지 스타일링 */
 .chatbox .messages .message.user {
     background: linear-gradient(145deg, #e1f5fe, #bbdefb);
@@ -190,7 +141,6 @@ h3 {
     transform: translateZ(10px);
     animation: messageIn 0.3s ease-out;
 }
-
 .chatbox .messages .message.bot {
     background: linear-gradient(145deg, #f5f5f5, #eeeeee);
     border-radius: 15px;
@@ -200,7 +150,6 @@ h3 {
     transform: translateZ(10px);
     animation: messageIn 0.3s ease-out;
 }
-
 /* 버튼 스타일링 */
 .duplicate-button {
     background: linear-gradient(145deg, #24292e, #1a1e22) !important;
@@ -212,12 +161,10 @@ h3 {
     border: none !important;
     cursor: pointer !important;
 }
-
 .duplicate-button:hover {
     transform: translateY(-2px) !important;
     box-shadow: 0 5px 15px rgba(0, 0, 0, 0.3) !important;
 }
-
 /* 입력 필드 스타일링 */
 """
 
@@ -228,13 +175,13 @@ with gr.Blocks(css=CSS) as demo:
         chatbot=chatbot,
         fill_height=True,
         theme="soft",
-        additional_inputs_accordion=gr.Accordion(label="⚙️ 옵션", open=False, render=False),
+        additional_inputs_accordion=gr.Accordion(label="⚙️ 옵션션", open=False, render=False),
         additional_inputs=[
             gr.Slider(
                 minimum=0,
                 maximum=1,
                 step=0.1,
-                value=0.3,
+                value=0.8,
                 label="온도",
                 render=False,
             ),