Update app.py

app.py

@@ -13,16 +13,13 @@ import torch
 from loguru import logger
 from PIL import Image
 from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer
-
-# CSV/TXT 분석
 import pandas as pd
-# PDF 텍스트 추출용
 import PyPDF2
 
 ##################################################
-# 상수 및 모델 로딩
+# 기본 설정
 ##################################################
-MAX_CONTENT_CHARS = 8000  # 텍스트로 전달 시 최대 8000자까지만
+MAX_CONTENT_CHARS = 8000  # 텍스트로 전달 시 최대 글자 수
 model_id = os.getenv("MODEL_ID", "google/gemma-3-27b-it")
 
 processor = AutoProcessor.from_pretrained(model_id, padding_side="left")
@@ -32,18 +29,17 @@ model = Gemma3ForConditionalGeneration.from_pretrained(
     torch_dtype=torch.bfloat16,
     attn_implementation="eager"
 )
-
 MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
 
-
 ##################################################
-# 1) CSV, TXT, PDF 분석 함수
+# 1) CSV, TXT, PDF 분석 함수 (빈 파일 대비)
 ##################################################
 def analyze_csv_file(path: str) -> str:
-    """CSV 파일 -> 문자열. 길면 잘라냄."""
     try:
         df = pd.read_csv(path)
-        df_str = df.to_string()
+        df_str = df.to_string().strip()
+        if not df_str:
+            df_str = "(CSV is empty)"
         if len(df_str) > MAX_CONTENT_CHARS:
             df_str = df_str[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
         return f"**[CSV File: {os.path.basename(path)}]**\n\n{df_str}"
@@ -52,10 +48,11 @@ def analyze_csv_file(path: str) -> str:
 
 
 def analyze_txt_file(path: str) -> str:
-    """TXT 파일 -> 전체 문자열. 길면 잘라냄."""
     try:
         with open(path, "r", encoding="utf-8") as f:
-            text = f.read()
+            text = f.read().strip()
+        if not text:
+            text = "(TXT is empty)"
         if len(text) > MAX_CONTENT_CHARS:
             text = text[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
         return f"**[TXT File: {os.path.basename(path)}]**\n\n{text}"
@@ -64,28 +61,26 @@ def analyze_txt_file(path: str) -> str:
 
 
 def pdf_to_markdown(pdf_path: str) -> str:
-    """PDF -> 텍스트 추출 -> Markdown. 길면 잘라냄."""
     try:
-        text_chunks = []
         with open(pdf_path, "rb") as f:
             reader = PyPDF2.PdfReader(f)
+            chunks = []
             for page_num, page in enumerate(reader.pages, start=1):
-                page_text = page.extract_text() or ""
-                page_text = page_text.strip()
-                if page_text:
-                    text_chunks.append(f"## Page {page_num}\n\n{page_text}\n")
+                ptext = (page.extract_text() or "").strip()
+                if ptext:
+                    chunks.append(f"## Page {page_num}\n\n{ptext}\n")
+        full_text = "\n".join(chunks).strip()
+        if not full_text:
+            full_text = "(PDF is empty)"
+        if len(full_text) > MAX_CONTENT_CHARS:
+            full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
+        return f"**[PDF File: {os.path.basename(pdf_path)}]**\n\n{full_text}"
     except Exception as e:
         return f"Failed to read PDF ({os.path.basename(pdf_path)}): {str(e)}"
 
-    full_text = "\n".join(text_chunks)
-    if len(full_text) > MAX_CONTENT_CHARS:
-        full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
-
-    return f"**[PDF File: {os.path.basename(pdf_path)}]**\n\n{full_text}"
-
 
 ##################################################
-# 2) 이미지/비디오 제한 검사 (CSV, PDF, TXT 제외)
+# 2) 이미지/비디오 업로드 제한
 ##################################################
 def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
     image_count = 0
@@ -102,9 +97,9 @@ def count_files_in_history(history: list[dict]) -> tuple[int, int]:
     image_count = 0
     video_count = 0
     for item in history:
+        # assistant 또는 content가 str이면 제외
         if item["role"] != "user" or isinstance(item["content"], str):
             continue
-        # item["content"]가 ["경로"] 형태일 때, 확장자를 확인
         file_path = item["content"][0]
         if file_path.endswith(".mp4"):
             video_count += 1
@@ -115,11 +110,10 @@ def count_files_in_history(history: list[dict]) -> tuple[int, int]:
 
 def validate_media_constraints(message: dict, history: list[dict]) -> bool:
     """
-    이미지 & 비디오 제한
+    이미지/비디오 개수 제한
     """
     media_files = []
     for f in message["files"]:
-        # 이미지/비디오만
         if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE) or f.endswith(".mp4"):
             media_files.append(f)
 
@@ -132,7 +126,7 @@ def validate_media_constraints(message: dict, history: list[dict]) -> bool:
     if video_count > 1:
         gr.Warning("Only one video is supported.")
         return False
-    # 비디오 + 이미지 혼합 불가
+    # 비디오+이미지 혼합 불가
     if video_count == 1:
         if image_count > 0:
             gr.Warning("Mixing images and videos is not allowed.")
@@ -144,7 +138,7 @@ def validate_media_constraints(message: dict, history: list[dict]) -> bool:
     if video_count == 0 and image_count > MAX_NUM_IMAGES:
         gr.Warning(f"You can upload up to {MAX_NUM_IMAGES} images.")
         return False
-    # <image> 태그와 실제 이미지 수 일치?
+    # <image> 태그 수와 이미지 파일 수 일치
     if "<image>" in message["text"] and message["text"].count("<image>") != new_image_count:
         gr.Warning("The number of <image> tags in the text does not match the number of images.")
         return False
@@ -182,7 +176,6 @@ def process_video(video_path: str) -> list[dict]:
             pil_image.save(temp_file.name)
             content.append({"type": "text", "text": f"Frame {timestamp}:"})
             content.append({"type": "image", "url": temp_file.name})
-    logger.debug(f"{content=}")
     return content
 
 
@@ -206,46 +199,51 @@ def process_interleaved_images(message: dict) -> list[dict]:
 
 
 ##################################################
-# 5) CSV/PDF/TXT는 텍스트 변환만, 이미지/비디오는 경로로
+# 5) CSV/PDF/TXT = 텍스트 / 이미지,비디오 = 실제 경로
 ##################################################
 def process_new_user_message(message: dict) -> list[dict]:
+    user_text = (message["text"] or "").strip() or "(No text)"
     if not message["files"]:
-        return [{"type": "text", "text": message["text"]}]
+        return [{"type": "text", "text": user_text}]
 
-    # 확장자별 분류
     video_files = [f for f in message["files"] if f.endswith(".mp4")]
     image_files = [f for f in message["files"] if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)]
     csv_files = [f for f in message["files"] if f.lower().endswith(".csv")]
     txt_files = [f for f in message["files"] if f.lower().endswith(".txt")]
     pdf_files = [f for f in message["files"] if f.lower().endswith(".pdf")]
 
-    # user 텍스트 추가
-    content_list = [{"type": "text", "text": message["text"]}]
+    content_list = [{"type": "text", "text": user_text}]
 
     # CSV
     for csv_path in csv_files:
         csv_analysis = analyze_csv_file(csv_path)
+        if not csv_analysis.strip():
+            csv_analysis = "(No CSV content?)"
         content_list.append({"type": "text", "text": csv_analysis})
 
     # TXT
     for txt_path in txt_files:
         txt_analysis = analyze_txt_file(txt_path)
+        if not txt_analysis.strip():
+            txt_analysis = "(No TXT content?)"
         content_list.append({"type": "text", "text": txt_analysis})
 
     # PDF
     for pdf_path in pdf_files:
-        pdf_markdown = pdf_to_markdown(pdf_path)
-        content_list.append({"type": "text", "text": pdf_markdown})
+        pdf_md = pdf_to_markdown(pdf_path)
+        if not pdf_md.strip():
+            pdf_md = "(No PDF content?)"
+        content_list.append({"type": "text", "text": pdf_md})
 
-    # 비디오
     if video_files:
+        # 하나만 처리
         content_list += process_video(video_files[0])
         return content_list
 
-    # 이미지
-    if "<image>" in message["text"]:
+    if "<image>" in user_text:
         return process_interleaved_images(message)
     else:
+        # 일반 이미지
         for img_path in image_files:
             content_list.append({"type": "image", "url": img_path})
 
@@ -253,13 +251,9 @@ def process_new_user_message(message: dict) -> list[dict]:
 
 
 ##################################################
-# 6) 히스토리 -> LLM 메시지 변환
+# 6) 히스토리 -> LLM 메시지 변환 (비이미지 경로는 무시)
 ##################################################
 def process_history(history: list[dict]) -> list[dict]:
-    """
-    여기서, 이미지/비디오 외의 파일(.csv, .pdf, .txt) 경로는
-    모델로 전달되지 않도록 제거 (or 무시)
-    """
     messages = []
     current_user_content = []
     for item in history:
@@ -267,71 +261,84 @@ def process_history(history: list[dict]) -> list[dict]:
             if current_user_content:
                 messages.append({"role": "user", "content": current_user_content})
                 current_user_content = []
-            # assistant -> 그냥 텍스트로
             messages.append({"role": "assistant", "content": [{"type": "text", "text": item["content"]}]})
         else:
             # user
             content = item["content"]
             if isinstance(content, str):
-                # 단순 텍스트
                 current_user_content.append({"type": "text", "text": content})
             else:
-                # 보통 [파일경로] 형태
-                file_path = content[0]
-                # 만약 이미지나 mp4가 아니라면 -> 무시
-                if re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE) or file_path.endswith(".mp4"):
-                    current_user_content.append({"type": "image", "url": file_path})
+                # [파일경로]
+                fpath = content[0]
+                # 이미지나 mp4만 유지, 나머지는 제외
+                if re.search(r"\.(png|jpg|jpeg|gif|webp)$", fpath, re.IGNORECASE) or fpath.endswith(".mp4"):
+                    current_user_content.append({"type": "image", "url": fpath})
                 else:
-                    # csv, pdf, txt 등은 제거
                     pass
     return messages
 
 
 ##################################################
-# 7) 메인 추론
+# 7) 메인 추론 (빈 토큰 방어)
 ##################################################
 @spaces.GPU(duration=120)
 def run(message: dict, history: list[dict], system_prompt: str = "", max_new_tokens: int = 512) -> Iterator[str]:
-    # a) 미디어 제한 검사
     if not validate_media_constraints(message, history):
         yield ""
         return
 
-    # b) 기존 히스토리 -> LLM 메시지
     messages = []
     if system_prompt:
         messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
     messages.extend(process_history(history))
-    messages.append({"role": "user", "content": process_new_user_message(message)})
 
-    # c) 모델 호출
-    inputs = processor.apply_chat_template(
+    user_content = process_new_user_message(message)
+    messages.append({"role": "user", "content": user_content})
+
+    # 1) tokenize=False 후 토큰 길이 체크
+    raw_text = processor.tokenizer.apply_chat_template(
         messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    token_ids = processor.tokenizer.encode(raw_text, add_special_tokens=False)
+    if len(token_ids) == 0:
+        # 빈 입력 → 임의 문구 추가
+        raw_text += " (No content?)"
+        token_ids = processor.tokenizer.encode(raw_text, add_special_tokens=False)
+
+    # 2) 실제 tokenizer
+    inputs = processor.tokenizer(
+        raw_text,
         return_tensors="pt",
-    ).to(device=model.device, dtype=torch.bfloat16)
+        padding=True
+    )
+    inputs = {k: v.to(model.device, dtype=torch.bfloat16) for k, v in inputs.items()}
 
-    streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
+    # 3) 스트리밍 생성
+    streamer = TextIteratorStreamer(processor.tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
     gen_kwargs = {
-        "inputs": inputs,
+        "inputs": inputs["input_ids"],
+        "attention_mask": inputs.get("attention_mask"),
         "streamer": streamer,
         "max_new_tokens": max_new_tokens,
+        "do_sample": True,
+        "temperature": 0.3,
+        "top_p": 0.95,
     }
+    gen_kwargs = {k: v for k, v in gen_kwargs.items() if v is not None}
+
     t = Thread(target=model.generate, kwargs=gen_kwargs)
     t.start()
 
     output = ""
-    for new_text in streamer:
-        output += new_text
+    for chunk in streamer:
+        output += chunk
         yield output
 
 
-
-
 ##################################################
-# 예시들 (한글화 버전)
+# 8) 예시
 ##################################################
 examples = [
 
@@ -463,17 +470,10 @@ examples = [
 ]
 
 
-
-
-
-##################################################
-# 9) Gradio ChatInterface
-##################################################
 demo = gr.ChatInterface(
     fn=run,
     type="messages",
     chatbot=gr.Chatbot(type="messages", scale=1, allow_tags=["image"]),
-    # 이미지(여러 확장자), mp4, csv, txt, pdf 허용
     textbox=gr.MultimodalTextbox(
         file_types=[
             ".png", ".jpg", ".jpeg", ".gif", ".webp",
@@ -488,13 +488,7 @@ demo = gr.ChatInterface(
             label="System Prompt",
             value="You are a deeply thoughtful AI. Consider problems thoroughly and derive correct solutions through systematic reasoning. Please answer in korean."
         ),
-        gr.Slider(
-            label="Max New Tokens",
-            minimum=100,
-            maximum=8000,
-            step=50,
-            value=2000
-        ),
+        gr.Slider(label="Max New Tokens", minimum=100, maximum=8000, step=50, value=2000),
     ],
     stop_btn=False,
     title="Gemma 3 27B IT",
@@ -505,7 +499,5 @@ demo = gr.ChatInterface(
     delete_cache=(1800, 1800),
 )
 
-
 if __name__ == "__main__":
     demo.launch()
-