Update app.py

app.py

@@ -188,9 +188,8 @@ def search_protein_data(analysis, dataset):
                 print(f"Error processing entry: {str(e)}")
                 continue
         
-        # 결과 정렬 및 반환
-        scored_entries.sort(reverse=True)
-        return scored_entries[:3]
+        # 결과 정렬 및 반환 (key 함수 사용)
+        return sorted(scored_entries, key=lambda x: x[0], reverse=True)[:3]
         
     except Exception as e:
         print(f"데이터 검색 중 오류: {str(e)}")
@@ -944,8 +943,25 @@ def process_chat_and_generate(message, history):
 
         # 2. 프롬프트 분석
         analysis = analyze_prompt(message)
+        if not analysis:
+            return history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": "요청 분석에 실패했습니다."}
+            ], None, None, None, None, None, None
+
         similar_structures = search_protein_data(analysis, ds)
+        if not similar_structures:
+            return history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": "적합한 참조 구조를 찾지 못했습니다."}
+            ], None, None, None, None, None, None
+
         params = extract_parameters(analysis, similar_structures)
+        if not params:
+            return history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": "파라미터 설정에 실패했습니다."}
+            ], None, None, None, None, None, None
 
         # 3. 분석 결과 추가 (이전 메시지 유지)
         current_history = current_history[:-1] + [
@@ -959,71 +975,115 @@ def process_chat_and_generate(message, history):
         yield (current_history, None, None, None, None, None, None)
 
         # 4. 단백질 생성
-        generator = protein_diffusion_model(...)  # 기존 파라미터 유지
-
-        # 5. 생성 과정 추적 (이전 메시지들 유지)
-        step = 0
-        for result in generator:
-            step += 1
-            progress_msg = f"단백질 생성 중... {step}/25 단계 완료"
-            current_history = current_history[:-1] + [
-                {"role": "assistant", "content": progress_msg}
-            ]
-            yield (
-                current_history,
-                create_radar_chart(calculate_hero_stats(
-                    params['helix_bias'],
-                    params['strand_bias'],
-                    params['loop_bias'],
-                    float(params['hydrophobic_target_score'])
-                )),
-                progress_msg,
-                result[0],
-                result[1],
-                result[2],
-                result[3]
+        try:
+            generator = protein_diffusion_model(
+                sequence=None,
+                seq_len=params['sequence_length'],
+                helix_bias=params['helix_bias'],
+                strand_bias=params['strand_bias'],
+                loop_bias=params['loop_bias'],
+                secondary_structure=None,
+                aa_bias=None,
+                aa_bias_potential=None,
+                num_steps="25",
+                noise="normal",
+                hydrophobic_target_score=str(params['hydrophobic_target_score']),
+                hydrophobic_potential="2",
+                contigs=None,
+                pssm=None,
+                seq_mask=None,
+                str_mask=None,
+                rewrite_pdb=None
             )
 
-        # 6. 최종 결과 및 설명 추가 (모든 이전 메시지 유지)
-        final_explanation = f"""
-        단백질 설계가 완료되었습니다.
+            # 5. 생성 과정 추적 (이전 메시지들 유지)
+            step = 0
+            final_result = None
+            for result in generator:
+                step += 1
+                final_result = result
+                progress_msg = f"단백질 생성 중... {step}/25 단계 완료"
+                current_history = current_history[:-1] + [
+                    {"role": "assistant", "content": progress_msg}
+                ]
+                yield (
+                    current_history,
+                    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
+                )
 
-        [분석 결과]
-        {analysis}
+            if final_result is None:
+                raise Exception("생성 결과가 없습니다")
 
-        [구조적 특징]
-        - 길이: {params['sequence_length']} 아미노산
-        - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
-        - 베타 시트 비율: {params['strand_bias']*100:.1f}%
-        - 루프 구조 비율: {params['loop_bias']*100:.1f}%
-        - 소수성 점수: {params['hydrophobic_target_score']}
+            # 6. 최종 결과 및 설명 추가
+            output_seq, output_pdb, structure_view, plddt_plot = final_result
+            
+            final_explanation = f"""
+            단백질 설계가 완료되었습니다.
 
-        [생성 과정]
-        - 총 {step}단계의 최적화 완료
-        - 최종 안정성 점수: {np.mean(plddt_data) if plddt_data else 0:.2f}
-        - 참조된 유사 구조: {len(similar_structures)}개
+            [분석 결과]
+            {analysis}
 
-        3D 구조와 상세 분석 결과를 확인하실 수 있습니다.
-        """
+            [구조적 특징]
+            - 길이: {params['sequence_length']} 아미노산
+            - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
+            - 베타 시트 비율: {params['strand_bias']*100:.1f}%
+            - 루프 구조 비율: {params['loop_bias']*100:.1f}%
+            - 소수성 점수: {params['hydrophobic_target_score']}
 
-        final_history = current_history + [
-            {"role": "assistant", "content": final_explanation}
-        ]
+            [생성 과정]
+            - 총 {step}단계의 최적화 완료
+            - 최종 안정성 점수: {np.mean(plddt_data) if plddt_data else 0:.2f}
+            - 참조된 유사 구조: {len(similar_structures)}개
 
-        return (
-            final_history,  # 모든 대화 기록 유지
-            create_radar_chart(stats),
-            final_explanation,
-            output_seq,
-            output_pdb,
-            structure_view,
-            plddt_plot
-        )
+            3D 구조와 상세 분석 결과를 확인하실 수 있습니다.
+            """
+
+            final_history = current_history + [
+                {"role": "assistant", "content": final_explanation}
+            ]
+
+            stats = calculate_hero_stats(
+                params['helix_bias'],
+                params['strand_bias'],
+                params['loop_bias'],
+                float(params['hydrophobic_target_score'])
+            )
+
+            return (
+                final_history,
+                create_radar_chart(stats),
+                final_explanation,
+                output_seq,
+                output_pdb,
+                structure_view,
+                plddt_plot
+            )
+
+        except Exception as e:
+            error_msg = f"단백질 생성 중 오류가 발생했습니다: {str(e)}"
+            print(error_msg)
+            traceback.print_exc()
+            return (
+                current_history + [
+                    {"role": "assistant", "content": error_msg}
+                ],
+                None, None, None, None, None, None
+            )
 
     except Exception as e:
+        error_msg = f"처리 중 오류가 발생했습니다: {str(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},
@@ -1032,10 +1092,6 @@ def process_chat_and_generate(message, history):
             None, None, None, None, None, None
         )
 
-# 필요한 모듈 추가
-import json
-from Bio import SeqIO, Align
-from Bio.Seq import Seq
 
 # 시작 부분에 추가
 def extract_keywords(analysis):
@@ -1066,9 +1122,10 @@ def calculate_similarity(keywords, entry):
     """키워드와 데이터셋 항목 간의 유사도 계산"""
     try:
         score = 0
-        if not entry or not isinstance(entry, dict):
+        if not isinstance(entry, dict):
             return 0
             
+        # 안전한 접근을 위한 get 메서드 사용
         sequence = entry.get('sequence', '').lower()
         description = entry.get('description', '')
         
@@ -1090,8 +1147,13 @@ def calculate_similarity(keywords, entry):
         print(f"유사도 계산 중 오류: {str(e)}")
         return 0
 
-# plddt_data 변수 추가
+# 전역 변수 정의
 plddt_data = []
+stats = {}  # 능력치 저장용
+output_seq = None
+output_pdb = None
+structure_view = None
+plddt_plot = None
 
 # 예제 프롬프트 처리 함수 수정
 def use_example(example):
@@ -1665,4 +1727,4 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
 
 # 실행
 demo.queue()
-demo.launch(debug=True)    
+demo.launch(debug=True, share=True)