Update app.py

app.py

@@ -6,6 +6,11 @@ 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"
@@ -32,58 +37,97 @@ h3 {
 }
 """
 
-# 모델과 토크나이저 로드
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# 디바이스 설정
+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
 
-# 데이터셋 로드
-dataset = load_dataset("elyza/ELYZA-tasks-100")
-print(dataset)
+# 데이터셋 로드 with 에러 처리
+try:
+    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]
+    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 = []
 
+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):
-    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()
+    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})
 
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
+        input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(input_ids, return_tensors="pt").to(device)
 
-chatbot = gr.Chatbot(height=500)
+        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()
+
+chatbot = gr.Chatbot(height=500)                                        
 
 CSS = """
 /* 전체 페이지 스타일링 */
@@ -184,21 +228,21 @@ 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.8,
+                value=0.3,
                 label="온도",
                 render=False,
             ),
             gr.Slider(
                 minimum=128,
-                maximum=1000000,
+                maximum=8000,
                 step=1,
-                value=100000,
+                value=4000,
                 label="최대 토큰 수",
                 render=False,
             ),