Update app.py

app.py

@@ -1,4 +1,7 @@
 import os
+
+os.system("pip install git+https://github.com/shumingma/transformers.git")
+
 import threading
 import torch
 import torch._dynamo
@@ -12,55 +15,17 @@ from transformers import (
 import gradio as gr
 import spaces
 
-# 필요한 경우 Bitnet 지원을 위한 transformers 설치
-# Hugging Face Spaces에서는 Dockerfile 등을 통해 미리 설치하는 것이 더 일반적입니다.
-# 로컬에서 테스트 시에는 필요할 수 있습니다.
-# print("Installing required transformers branch...")
-# try:
-#     os.system("pip install git+https://github.com/shumingma/transformers.git -q")
-#     print("transformers branch installed.")
-# except Exception as e:
-#     print(f"Error installing transformers branch: {e}")
-#     print("Proceeding with potentially default transformers version.")
-
-# os.system("pip install accelerate bitsandbytes -q") # bitsandbytes, accelerate도 필요할 수 있습니다.
-
-
 model_id = "microsoft/bitnet-b1.58-2B-4T"
 
-# 모델 및 토크나이저 로드
-print(f"Loading model: {model_id}")
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    # device_map="auto"는 여러 GPU 또는 CPU로 모델을 자동으로 분산 로드합니다.
-    # bfloat16은 모델 가중치에 사용되는 데이터 타입입니다.
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        # load_in_8bit=True # Bitnet은 1.58bit이므로 8bit 로딩이 의미 없을 수 있습니다.
-    )
-    print(f"Model loaded successfully on device: {model.device}")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    # 모델 로딩 실패 시 더미 모델 사용 또는 오류 처리
-    class DummyModel:
-        def generate(self, **kwargs):
-            # 더미 응답 생성
-            input_ids = kwargs.get('input_ids')
-            streamer = kwargs.get('streamer')
-            if streamer:
-                # 간단한 더미 응답 스트리밍
-                dummy_response = "모델 로딩에 실패하여 더미 응답을 제공합니다. 설정/경로를 확인하세요."
-                for char in dummy_response:
-                    streamer.put(char)
-                streamer.end()
-    model = DummyModel()
-    tokenizer = AutoTokenizer.from_pretrained("gpt2") # 더미 토크나이저
-    print("Using dummy model due to loading failure.")
-
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+print(model.device)
 
-@spaces.GPU # Hugging Face Spaces에서 GPU 사용을 지정합니다.
+@spaces.GPU
 def respond(
     message: str,
     history: list[tuple[str, str]],
@@ -81,11 +46,6 @@ def respond(
     Yields:
         The growing response text as new tokens are generated.
     """
-    # 더미 모델 사용 시 스트리밍 오류 방지
-    if isinstance(model, DummyModel):
-         yield "모델 로딩에 실패하여 응답을 생성할 수 없습니다."
-         return
-
     messages = [{"role": "system", "content": system_message}]
     for user_msg, bot_msg in history:
         if user_msg:
@@ -94,236 +54,45 @@ def respond(
             messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    try:
-        prompt = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-        streamer = TextIteratorStreamer(
-            tokenizer, skip_prompt=True, skip_special_tokens=True
-        )
-        generate_kwargs = dict(
-            **inputs,
-            streamer=streamer,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            # Bitnet 모델에 필요한 추가 인자 설정 (모델 문서 확인 필요)
-            # 예를 들어, quantize_config 등
-        )
-
-        # 쓰레드에서 모델 생성 실행
-        thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
-
-        # 스트리머로부터 텍스트를 읽어와 yield
-        response = ""
-        for new_text in streamer:
-            # yield 하기 전에 불필요한 공백/토큰 제거 또는 처리 가능
-            response += new_text
-            yield response
-
-    except Exception as e:
-        print(f"Error during response generation: {e}")
-        yield f"응답 생성 중 오류가 발생했습니다: {e}"
-
-
-# --- 디자인 개선을 위한 CSS 코드 ---
-css_styles = """
-/* 전체 페이지 배경 및 기본 폰트 설정 */
-body {
-    font-family: 'Segoe UI', 'Roboto', 'Arial', sans-serif;
-    line-height: 1.6;
-    margin: 0;
-    padding: 20px; /* 앱 주변 여백 추가 */
-    background-color: #f4f7f6; /* 부드러운 배경색 */
-}
-
-/* 메인 앱 컨테이너 스타일 */
-.gradio-container {
-    max-width: 900px; /* 중앙 정렬 및 최대 너비 제한 */
-    margin: 20px auto;
-    border-radius: 12px; /* 둥근 모서리 */
-    overflow: hidden; /* 자식 요소들이 모서리를 넘지 않도록 */
-    box-shadow: 0 8px 16px rgba(0, 0, 0, 0.1); /* 그림자 효과 */
-    background-color: #ffffff; /* 앱 내용 영역 배경색 */
-}
-
-/* 타이틀 및 설명 영역 (ChatInterface의 기본 타이틀/설명) */
-/* 이 영역은 ChatInterface 구조에 따라 정확한 클래스 이름이 다를 수 있으나,
-   .gradio-container 내부의 첫 블록이나 H1/P 태그를 타겟할 수 있습니다.
-   테마와 함께 사용하면 대부분 잘 처리됩니다. 여기서는 추가적인 패딩 등만 고려 */
-.gradio-container > .gradio-block:first-child {
-    padding: 20px 20px 10px 20px; /* 상단 패딩 조정 */
-}
-
-/* 채팅 박스 영역 스타일 */
-.gradio-chatbox {
-    /* 테마에 의해 스타일링되지만, 추가적인 내부 패딩 등 조정 가능 */
-     padding: 15px;
-     background-color: #fefefe; /* 채팅 영역 배경색 */
-     border-radius: 8px; /* 채팅 영역 내부 모서리 */
-     border: 1px solid #e0e0e0; /* 경계선 */
-}
-
-/* 채팅 메시지 스타일 */
-.gradio-chatmessage {
-    margin-bottom: 12px;
-    padding: 10px 15px;
-    border-radius: 20px; /* 둥근 메시지 모서리 */
-    max-width: 75%; /* 메시지 너비 제한 */
-    word-wrap: break-word; /* 긴 단어 줄바꿈 */
-    white-space: pre-wrap; /* 공백 및 줄바꿈 유지 */
-    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05); /* 메시지에 약간의 그림자 */
-}
-
-/* 사용자 메시지 스타일 */
-.gradio-chatmessage.user {
-    background-color: #007bff; /* 파란색 계열 */
-    color: white;
-    margin-left: auto; /* 오른쪽 정렬 */
-    border-bottom-right-radius: 2px; /* 오른쪽 아래 모서리 각지게 */
-}
-
-/* 봇 메시지 스타일 */
-.gradio-chatmessage.bot {
-    background-color: #e9ecef; /* 밝은 회색 */
-    color: #333; /* 어두운 텍스트 */
-    margin-right: auto; /* 왼쪽 정렬 */
-    border-bottom-left-radius: 2px; /* 왼쪽 아래 모서리 각지게 */
-}
-
-/* 입력창 및 버튼 영역 스타일 */
-.gradio-input-box {
-     padding: 15px;
-     border-top: 1px solid #eee; /* 위쪽 경계선 */
-     background-color: #f8f9fa; /* 입력 영역 배경색 */
-}
- /* 입력 텍스트 에어리어 스타일 */
- .gradio-input-box textarea {
-     border-radius: 8px;
-     padding: 10px;
-     border: 1px solid #ccc;
-     resize: none !important; /* 입력창 크기 조절 비활성화 (선택 사항) */
-     min-height: 50px; /* 최소 높이 */
-     max-height: 150px; /* 최대 높이 */
-     overflow-y: auto; /* 내용 넘칠 경우 스크롤 */
- }
- /* 스크롤바 스타일 (선택 사항) */
- .gradio-input-box textarea::-webkit-scrollbar {
-    width: 8px;
- }
- .gradio-input-box textarea::-webkit-scrollbar-thumb {
-    background-color: #ccc;
-    border-radius: 4px;
- }
- .gradio-input-box textarea::-webkit-scrollbar-track {
-    background-color: #f1f1f1;
- }
-
-
-/* 버튼 스타일 */
-.gradio-button {
-    border-radius: 8px;
-    padding: 10px 20px;
-    font-weight: bold;
-    transition: background-color 0.2s ease, opacity 0.2s ease; /* 호버 애니메이션 */
-    border: none; /* 기본 테두리 제거 */
-    cursor: pointer;
-}
-
-.gradio-button:not(.clear-button) { /* Send 버튼 */
-    background-color: #28a745; /* 초록색 */
-    color: white;
-}
-.gradio-button:not(.clear-button):hover {
-    background-color: #218838;
-}
-.gradio-button:disabled { /* 비활성화된 버튼 */
-    opacity: 0.6;
-    cursor: not-allowed;
-}
-
-
-.gradio-button.clear-button { /* Clear 버튼 */
-    background-color: #dc3545; /* 빨간색 */
-    color: white;
-}
-.gradio-button.clear-button:hover {
-    background-color: #c82333;
-}
-
-/* Additional inputs (추가 설정) 영역 스타일 */
-/* 이 영역은 보통 아코디언 형태로 되어 있으며, .gradio-accordion 클래스를 가집니다. */
- .gradio-accordion {
-     border-radius: 12px; /* 외부 컨테이너와 동일한 모서리 */
-     margin-top: 15px; /* 채팅 영역과의 간격 */
-     border: 1px solid #ddd; /* 경계선 */
-     box-shadow: none; /* 내부 그림자 제거 */
- }
- /* 아코디언 헤더 (라벨) 스타일 */
- .gradio-accordion .label {
-      font-weight: bold;
-      color: #007bff; /* 파란색 계열 */
-      padding: 15px; /* 헤더 패딩 */
-      background-color: #e9ecef; /* 헤더 배경색 */
-      border-bottom: 1px solid #ddd; /* 헤더 아래 경계선 */
-      border-top-left-radius: 11px; /* 상단 모서리 */
-      border-top-right-radius: 11px;
- }
-  /* 아코디언 내용 영역 스타일 */
- .gradio-accordion .wrap {
-      padding: 15px; /* 내용 패딩 */
-      background-color: #fefefe; /* 내용 배경색 */
-      border-bottom-left-radius: 11px; /* 하단 모서리 */
-      border-bottom-right-radius: 11px;
- }
-/* 추가 설정 내 개별 입력 컴포넌트 스���일 (슬라이더, 텍스트박스 등) */
-.gradio-slider, .gradio-textbox, .gradio-number {
-     margin-bottom: 10px; /* 각 입력 요소 아래 간격 */
-     padding: 8px; /* 내부 패딩 */
-     border: 1px solid #e0e0e0; /* 경계선 */
-     border-radius: 8px; /* 둥근 모서리 */
-     background-color: #fff; /* 배경색 */
-}
-/* 입력 필드 라벨 스타일 */
- .gradio-label {
-     font-weight: normal; /* 라벨 폰트 굵기 */
-     margin-bottom: 5px; /* 라벨과 입력 필드 간 간격 */
-     color: #555; /* 라벨 색상 */
-     display: block; /* 라벨을 블록 요소로 만들어 위로 올림 */
- }
-/* 슬라이더 트랙 및 핸들 스타일 (더 세밀한 조정 가능) */
-/* 예: .gradio-slider input[type="range"]::-webkit-slider-thumb {} */
-
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-/* 마크다운/HTML 컴포넌트 내 스타일 */
-.gradio-markdown, .gradio-html {
-    padding: 10px 0; /* 상하 패딩 */
-}
-"""
-# --- 디자인 개선을 위한 CSS 코드 끝 ---
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
 
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response
 
-# Gradio 인터페이스 설정
 demo = gr.ChatInterface(
     fn=respond,
-    # 타이틀 및 설명에 HTML 태그 사용 예시 (<br> 태그 사용)
-    title="<h1 style='text-align: center; color: #007bff;'>Bitnet-b1.58-2B-4T Chatbot</h1>",
-    description="<p style='text-align: center; color: #555;'>This chat application is powered by Microsoft's SOTA Bitnet-b1.58-2B-4T and designed for natural and fast conversations.</p>",
+    title="Bitnet-b1.58-2B-4T",
+    description="Bitnet-b1.58-2B-4T",
     examples=[
         [
-            "Hello! How are you?",
-            "You are a helpful AI assistant for everyday tasks.",
+            "Hello!",
+            "You are a helpful AI.",
             512,
             0.7,
             0.95,
         ],
         [
-            "Can you code a snake game in Python?",
-            "You are a helpful AI assistant for coding.",
+            "Can you code a snake game?",
+            "You are a helpful AI.",
             2048,
             0.7,
             0.95,
@@ -332,8 +101,7 @@ demo = gr.ChatInterface(
     additional_inputs=[
         gr.Textbox(
             value="You are a helpful AI assistant.",
-            label="System message",
-            lines=3 # 시스템 메시지 입력창 높이 조절
+            label="System message"
         ),
         gr.Slider(
             minimum=1,
@@ -357,14 +125,7 @@ demo = gr.ChatInterface(
             label="Top-p (nucleus sampling)"
         ),
     ],
-    # 테마 적용 (여러 테마 중 선택 가능: gr.themes.Soft(), gr.themes.Glass(), gr.themes.Default(), etc.)
-    theme=gr.themes.Soft(),
-    # 커스텀 CSS 적용
-    css=css_styles,
 )
 
-# 애플리케이션 실행
 if __name__ == "__main__":
-    # launch(share=True)는 퍼블릭 URL 생성 (디버깅/공유 목적, 주의 필요)
-    demo.launch()
-    # demo.launch(debug=True) # 디버깅 모드 활성화
+    demo.launch()