Update app.py

app.py

@@ -68,13 +68,17 @@ client = OpenAI(
     api_key=ACCESS_TOKEN,
 )
 
-# 데이터셋 로드
-ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
-                 token=ACCESS_TOKEN)
-
-current_protein_result = None
-
 
+# 데이터셋 로드 및 구조 확인
+try:
+    ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
+                     token=ACCESS_TOKEN)
+    print("Dataset structure:", ds)
+    print("First entry example:", next(iter(ds['train'])))
+except Exception as e:
+    print(f"Dataset loading error: {str(e)}")
+    raise
+    
 def respond(
     message,
     history,
@@ -144,22 +148,33 @@ def analyze_prompt(message):
 def search_protein_data(analysis, dataset):
     """분석 결과를 바탕으로 데이터셋에서 유사한 구조 검색"""
     try:
-        # 분석에서 키워드 추출
+        # 키워드 추출
         keywords = extract_keywords(analysis)
+        print("Extracted keywords:", keywords)
         
+        # 데이터셋 구조 확인
+        if not dataset or 'train' not in dataset:
+            print("Invalid dataset structure")
+            return []
+            
         # 유사도 점수 계산
         scored_entries = []
         for entry in dataset['train']:
-            score = calculate_similarity(keywords, entry)
-            scored_entries.append((score, entry))
+            try:
+                score = calculate_similarity(keywords, entry)
+                scored_entries.append((score, entry))
+            except Exception as e:
+                print(f"Error processing entry: {str(e)}")
+                continue
         
-        # 상위 유사 구조 반환
+        # 결과 정렬 및 반환
         scored_entries.sort(reverse=True)
-        return scored_entries[:3]  # 상위 3개 구조 반환
+        return scored_entries[:3]
+        
     except Exception as e:
         print(f"데이터 검색 중 오류: {str(e)}")
         return []
-
+        
 def extract_parameters(analysis, similar_structures):
     """분석 결과와 유사 구조를 바탕으로 생성 파라미터 결정"""
     try:
@@ -969,6 +984,40 @@ def extract_parameters_from_chat(chat_response):
         print(f"파라미터 추출 중 오류: {str(e)}")
         return None
 
+def download_checkpoint_files():
+    """필요한 체크포인트 파일 다운로드"""
+    try:
+        import requests
+        
+        # 체크포인트 파일 URL (실제 URL로 교체 필요)
+        dssp_url = "YOUR_DSSP_CHECKPOINT_URL"
+        og_url = "YOUR_OG_CHECKPOINT_URL"
+        
+        # DSSP 체크포인트 다운로드
+        if not os.path.exists(dssp_checkpoint):
+            print("Downloading DSSP checkpoint...")
+            response = requests.get(dssp_url)
+            with open(dssp_checkpoint, 'wb') as f:
+                f.write(response.content)
+        
+        # OG 체크포인트 다운로드
+        if not os.path.exists(og_checkpoint):
+            print("Downloading OG checkpoint...")
+            response = requests.get(og_url)
+            with open(og_checkpoint, 'wb') as f:
+                f.write(response.content)
+                
+        print("Checkpoint files downloaded successfully")
+    except Exception as e:
+        print(f"Error downloading checkpoint files: {str(e)}")
+        raise
+
+# 시작 시 체크포인트 파일 확인 및 다운로드
+try:
+    download_checkpoint_files()
+except Exception as e:
+    print(f"Warning: Could not download checkpoint files: {str(e)}")
+
 with gr.Blocks(theme='ParityError/Interstellar') as demo:
     with gr.Row():
         with gr.Column(scale=1):