Update app.py

app.py

@@ -909,12 +909,30 @@ def update_protein_display(chat_response):
             }
     return None
 
+def analyze_active_sites(sequence):
+    """활성 부위 분석"""
+    return "분석 중..."  # 임시 구현
+
+def predict_interactions(params):
+    """상호작용 예측"""
+    return "예측 중..."  # 임시 구현
+
+def evaluate_stability(plddt_data):
+    """안정성 평가"""
+    if not plddt_data:
+        return "평가 불가"
+    avg_score = np.mean(plddt_data)
+    if avg_score > 0.8:
+        return "매우 안정적"
+    elif avg_score > 0.6:
+        return "안정적"
+    else:
+        return "보통"
+
 def process_chat_and_generate(message, history):
     try:
-        # 초기 응답 생성
+        # 1. 초기 응답 생성
         initial_response = "단백질 설계를 시작합니다. 잠시만 기다려주세요..."
-        
-        # 초기 응답 표시
         yield (
             history + [
                 {"role": "user", "content": message},
@@ -923,13 +941,22 @@ def process_chat_and_generate(message, history):
             None, None, None, None, None, None
         )
 
-
-        # 1. 프롬프트 분석 및 파라미터 설정
+        # 2. 프롬프트 분석
         analysis = analyze_prompt(message)
         similar_structures = search_protein_data(analysis, ds)
         params = extract_parameters(analysis, similar_structures)
-        
-        # 2. 단백질 생성
+
+        # 3. 단백질 생성 시작 메시지
+        progress_msg = "단백질 구조 생성을 시작합니다..."
+        yield (
+            history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": progress_msg}
+            ],
+            None, None, None, None, None, None
+        )
+
+        # 4. 단백질 생성
         generator = protein_diffusion_model(
             sequence=None,
             seq_len=params['sequence_length'],
@@ -949,102 +976,95 @@ def process_chat_and_generate(message, history):
             str_mask=None,
             rewrite_pdb=None
         )
-        
-        # 3. 모든 단계의 결과 수집
+
+        # 5. 생성 과정 추적
         final_result = None
         plddt_data = []
-        
-        # 히어로 능력치 계산
-        stats = calculate_hero_stats(
-            helix_bias=params['helix_bias'],
-            strand_bias=params['strand_bias'],
-            loop_bias=params['loop_bias'],
-            hydrophobic_score=float(params['hydrophobic_target_score'])
-        )
-        
+        step = 0
+
         for result in generator:
+            step += 1
             final_result = result
             if result[3]:  # plddt_plot이 존재하는 경우
                 ax = result[3].gca()
                 if ax.lines:
                     line = ax.lines[0]
                     plddt_data = line.get_ydata().tolist()
-            
-            # 중간 결과 업데이트
+
+            progress_msg = f"단백질 생성 중... {step}/25 단계 완료"
             yield (
                 history + [
                     {"role": "user", "content": message},
-                    {"role": "assistant", "content": f"단백질 생성 중... {len(plddt_data)}단계 완료"}
+                    {"role": "assistant", "content": progress_msg}
                 ],
-                create_radar_chart(stats),
-                f"단백질 생성 진행 중... {len(plddt_data)}단계 완료",
+                create_radar_chart(calculate_hero_stats(
+                    params['helix_bias'],
+                    params['strand_bias'],
+                    params['loop_bias'],
+                    float(params['hydrophobic_target_score'])
+                )),
+                progress_msg,
                 result[0],  # output_seq
                 result[1],  # output_pdb
                 result[2],  # structure_view
                 result[3]   # plddt_plot
             )
-        
-        if final_result is None:
-            raise Exception("생성 결과가 없습니다")
-            
-        output_seq, output_pdb, structure_view, plddt_plot = final_result
-        
 
-        
+        # 6. 최종 결과 및 설명 생성
+        if final_result:
+            output_seq, output_pdb, structure_view, plddt_plot = final_result
+            
+            final_explanation = f"""
+            단백질 설계가 완료되었습니다.
 
-        # 생성 완료 후 상세 설명 추가
-        final_explanation = f"""
-        단백질 설계가 완료되었습니다.
+            [분석 결과]
+            {analysis}
 
-        [분석 결과]
-        {analysis}
+            [구조적 특징]
+            - 길이: {params['sequence_length']} 아미노산
+            - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
+            - 베타 시트 비율: {params['strand_bias']*100:.1f}%
+            - 루프 구조 비율: {params['loop_bias']*100:.1f}%
+            - 소수성 점수: {params['hydrophobic_target_score']}
 
-        [구조적 특징]
-        - 길이: {params['sequence_length']} 아미노산
-        - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
-        - 베타 시트 비율: {params['strand_bias']*100:.1f}%
-        - 루프 구조 비율: {params['loop_bias']*100:.1f}%
-        - 소수성 점수: {params['hydrophobic_target_score']}
+            [생성 과정]
+            - 총 {len(plddt_data)}단계의 최적화 완료
+            - 최종 안정성 점수: {np.mean(plddt_data) if plddt_data else 0:.2f}
+            - 참조된 유사 구조: {len(similar_structures)}개
 
-        [생성 과정]
-        - 총 {len(plddt_data)}단계의 최적화 수행
-        - 최종 안정성 점수: {np.mean(plddt_data) if plddt_data else 0:.2f}
-        - 참조된 유사 구조: {len(similar_structures)}개
+            3D 구조와 상세 분석 결과를 확인하실 수 있습니다.
+            """
 
-        [기능적 특성]
-        - 주요 활성 부위: {analyze_active_sites(output_seq)}
-        - 예상되는 상호작용: {predict_interactions(params)}
-        - 안정성 평가: {evaluate_stability(plddt_data)}
+            return (
+                history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": final_explanation}
+                ],
+                create_radar_chart(calculate_hero_stats(
+                    params['helix_bias'],
+                    params['strand_bias'],
+                    params['loop_bias'],
+                    float(params['hydrophobic_target_score'])
+                )),
+                final_explanation,
+                output_seq,
+                output_pdb,
+                structure_view,
+                plddt_plot
+            )
 
-        3D 구조와 상세 분석 결과를 확인하실 수 있습니다.
-        """
-        
-        return (
-            history + [
-                {"role": "user", "content": message},
-                {"role": "assistant", "content": final_explanation}
-            ],
-            create_radar_chart(stats),
-            final_explanation,
-            output_seq,
-            output_pdb,
-            structure_view,
-            plddt_plot
-        )
-        
     except Exception as e:
         print(f"Error in process_chat_and_generate: {str(e)}")
         traceback.print_exc()
+        error_msg = f"오류가 발생했습니다: {str(e)}"
         return (
             history + [
                 {"role": "user", "content": message},
-                {"role": "assistant", "content": f"오류가 발생했습니다: {str(e)}"}
+                {"role": "assistant", "content": error_msg}
             ],
             None, None, None, None, None, None
         )
 
-
-
 def extract_keywords(analysis):
     """분석 텍스트에서 키워드 추출"""
     try:
@@ -1567,7 +1587,8 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
               [chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
-    # 이벤트 핸들러 수정
+
+# 이벤트 핸들러 연결
     msg.submit(
         fn=process_chat_and_generate,
         inputs=[msg, chatbot],
@@ -1581,11 +1602,23 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
             plddt_plot
         ],
         show_progress=True
-    ).then(
-        lambda: None,  # 진행 상태 초기화
-        None,
-        None
     )
+    
+    submit_btn.click(
+        fn=process_chat_and_generate,
+        inputs=[msg, chatbot],
+        outputs=[
+            chatbot,
+            hero_stats,
+            hero_description,
+            output_seq,
+            output_pdb,
+            output_viewer,
+            plddt_plot
+        ],
+        show_progress=True
+    )
+
 
 # 실행
 demo.queue()