Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,11 +1,12 @@
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import gdown
 import zipfile
 import os
 import tempfile
 import json
+import urllib.request
+import requests
 
 # 페이지 설정
 st.set_page_config(
@@ -17,6 +18,114 @@ st.set_page_config(
 # 구글 드라이브 파일 ID
 GDRIVE_FILE_ID = "1RilOqw77G_Kr9_Cl6kR2UTuBa3uvDWQ8"
 
+def download_from_google_drive(file_id, destination):
+    """검색 결과 기반 구글 드라이브 다운로드 함수"""
+    try:
+        # 방법 1: confirm=t 파라미터 사용 (검색 결과 [3], [5]에서 확인)
+        st.info("📥 방법 1: confirm=t 파라미터로 다운로드 시도...")
+        url = f"https://drive.google.com/uc?export=download&id={file_id}&confirm=t"
+        
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(destination, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            return True
+        
+        # 방법 2: 쿠키를 사용한 다단계 다운로드 (검색 결과 [5]에서 확인)
+        st.info("📥 방법 2: 쿠키 기반 다운로드 시도...")
+        session = requests.Session()
+        
+        # 첫 번째 요청으로 쿠키 획득
+        response = session.get(f"https://drive.google.com/uc?export=download&id={file_id}")
+        
+        # 확인 토큰 추출
+        for key, value in response.cookies.items():
+            if key.startswith('download_warning'):
+                token = value
+                break
+        else:
+            token = 't'  # 기본값
+        
+        # 두 번째 요청으로 실제 다운로드
+        params = {'export': 'download', 'id': file_id, 'confirm': token}
+        response = session.get("https://drive.google.com/uc", params=params, stream=True)
+        
+        if response.status_code == 200:
+            with open(destination, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            return True
+        
+        # 방법 3: urllib 사용
+        st.info("📥 방법 3: urllib로 다운로드 시도...")
+        urllib.request.urlretrieve(
+            f"https://drive.google.com/uc?export=download&id={file_id}&confirm=t",
+            destination
+        )
+        return True
+        
+    except Exception as e:
+        st.error(f"다운로드 실패: {e}")
+        return False
+
+def try_gdown_download(file_id, destination):
+    """gdown을 사용한 다운로드 (검색 결과 해결책 적용)"""
+    try:
+        # 검색 결과 [4]에서 확인: 쿠키 파일 삭제
+        import gdown
+        
+        # 쿠키 파일 경로들
+        cookie_paths = [
+            os.path.expanduser("~/.cache/gdown/cookies.txt"),
+            os.path.expanduser("~/.cache/gdown/cookies.json"),
+        ]
+        
+        # 기존 쿠키 파일 삭제
+        for cookie_path in cookie_paths:
+            if os.path.exists(cookie_path):
+                os.remove(cookie_path)
+                st.info(f"🗑️ 기존 쿠키 파일 삭제: {cookie_path}")
+        
+        # 검색 결과 [4]에서 확인: gdown 4.6.0 버전 사용 권장
+        st.info("📥 gdown으로 다운로드 시도...")
+        
+        # 방법 1: ID만 사용
+        try:
+            gdown.download(id=file_id, output=destination, quiet=False)
+            return True
+        except:
+            pass
+        
+        # 방법 2: confirm=t 추가 (검색 결과 [6]에서 확인)
+        try:
+            gdown.download(f"{file_id}&confirm=t", output=destination, quiet=False)
+            return True
+        except:
+            pass
+        
+        # 방법 3: use_cookies=True 옵션 (검색 결과 [6]에서 확인)
+        try:
+            gdown.download(
+                f"https://drive.google.com/uc?id={file_id}",
+                output=destination,
+                quiet=False,
+                fuzzy=True,
+                use_cookies=True
+            )
+            return True
+        except:
+            pass
+        
+        return False
+        
+    except ImportError:
+        st.warning("gdown이 설치되지 않았습니다. 대안 방법을 사용합니다.")
+        return False
+    except Exception as e:
+        st.warning(f"gdown 다운로드 실패: {e}")
+        return False
+
 @st.cache_resource
 def download_and_load_model():
     """구글 드라이브에서 모델 다운로드 및 로드"""
@@ -26,42 +135,27 @@ def download_and_load_model():
         zip_path = os.path.join(temp_dir, "koalpaca_cpu_deployment.zip")
         extract_path = os.path.join(temp_dir, "koalpaca_cpu_deployment")
         
-        # 구글 드라이브에서 다운로드 (검색 결과 해결책 적용)
-        st.info("📥 모델 다운로드 중...")
+        st.info("📥 모델 다운로드 시작...")
         
-        # 방법 1: confirm=t 파라미터 추가 (대용량 파일���)
-        download_url = f"https://drive.google.com/uc?id={GDRIVE_FILE_ID}&confirm=t"
+        # 다운로드 시도 (여러 방법)
+        download_success = False
         
-        try:
-            # fuzzy=True와 use_cookies=True 옵션 추가
-            gdown.download(
-                download_url, 
-                zip_path, 
-                quiet=False,
-                fuzzy=True,
-                use_cookies=True
-            )
-        except Exception as e1:
-            st.warning(f"첫 번째 다운로드 방법 실패: {e1}")
-            try:
-                # 방법 2: 파일 ID만 사용
-                st.info("대안 다운로드 방법 시도 중...")
-                gdown.download(
-                    id=GDRIVE_FILE_ID,
-                    output=zip_path,
-                    quiet=False
-                )
-            except Exception as e2:
-                st.error(f"두 번째 다운로드 방법도 실패: {e2}")
-                # 방법 3: 직접 URL 구성
-                st.info("최종 다운로드 방법 시도 중...")
-                import urllib.request
-                direct_url = f"https://drive.google.com/uc?export=download&id={GDRIVE_FILE_ID}"
-                urllib.request.urlretrieve(direct_url, zip_path)
+        # 1. gdown 시도
+        if try_gdown_download(GDRIVE_FILE_ID, zip_path):
+            download_success = True
+            st.success("✅ gdown으로 다운로드 성공!")
+        
+        # 2. requests 기반 다운로드 시도
+        elif download_from_google_drive(GDRIVE_FILE_ID, zip_path):
+            download_success = True
+            st.success("✅ requests로 다운로드 성공!")
+        
+        if not download_success:
+            raise Exception("모든 다운로드 방법이 실패했습니다.")
         
         # 파일 다운로드 확인
-        if not os.path.exists(zip_path):
-            raise FileNotFoundError("모델 파일 다운로드에 실패했습니다.")
+        if not os.path.exists(zip_path) or os.path.getsize(zip_path) < 1024:
+            raise FileNotFoundError("다운로드된 파일이 유효하지 않습니다.")
         
         st.success(f"✅ 다운로드 완료: {os.path.getsize(zip_path) / 1024**3:.2f}GB")
         
@@ -188,6 +282,7 @@ def main():
         1. 구글 드라이브 파일 권한을 '링크가 있는 모든 사용자'로 설정
         2. 파일 ID가 정확한지 확인
         3. 파일 크기가 너무 큰 경우 시간이 오래 걸릴 수 있음
+        4. gdown 쿠키 문제는 자동으로 해결됩니다
         """)
         return