Update app.py

app.py

@@ -66,6 +66,9 @@ client = OpenAI(
 ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
                  token=ACCESS_TOKEN)
 
+current_protein_result = None
+
+
 def respond(
     message,
     history,
@@ -796,6 +799,157 @@ def combined_generation(name, strength, flexibility, speed, defense, size, abili
             None
         )
 
+
+def process_chat_and_generate(message, history):
+    try:
+        # 1. 프롬프트 분석 및 파라미터 설정
+        analysis = analyze_prompt(message)
+        similar_structures = search_protein_data(analysis, ds)
+        params = extract_parameters(analysis, similar_structures)
+        
+        # 2. 단백질 생성
+        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
+        )
+        
+        # 3. 생성 결과 가져오기
+        output_seq, output_pdb, structure_view, plddt_plot = next(generator)
+        
+        # 4. 설명 생성
+        explanation = f"""
+        요청하신 기능에 맞는 단백질을 생성했습니다:
+        
+        분석된 요구사항:
+        {analysis}
+        
+        설계된 구조적 특징:
+        - 길이: {params['sequence_length']} 아미노산
+        - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
+        - 베타 시트 비율: {params['strand_bias']*100:.1f}%
+        - 루프 구조 비율: {params['loop_bias']*100:.1f}%
+        - 소수성 점수: {params['hydrophobic_target_score']}
+        """
+        
+        # 5. 챗봇 히스토리 업데이트
+        new_history = history + [
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": explanation}
+        ]
+        
+        # 6. 모든 결과 반환
+        return (
+            new_history,  # 챗봇 히스토리
+            create_radar_chart(calculate_hero_stats(params)),  # 능력치 차트
+            explanation,  # 설명
+            output_seq,  # 단백질 서열
+            output_pdb,  # PDB 파일
+            structure_view,  # 3D 구조
+            plddt_plot  # 신뢰도 플롯
+        )
+        
+    except Exception as e:
+        print(f"Error in process_chat_and_generate: {str(e)}")
+        return (
+            history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": f"오류가 발생했습니다: {str(e)}"}
+            ],
+            None, None, None, None, None, None
+        )
+
+def extract_keywords(analysis):
+    """분석 텍스트에서 키워드 추출"""
+    try:
+        # 기본 키워드 추출
+        keywords = []
+        # 주요 기능 키워드
+        if "치료" in analysis: keywords.extend(["therapeutic", "binding"])
+        if "결합" in analysis: keywords.extend(["binding", "interaction"])
+        if "촉매" in analysis: keywords.extend(["enzyme", "catalytic"])
+        
+        # 환경 키워드
+        if "막" in analysis: keywords.extend(["membrane", "transmembrane"])
+        if "수용성" in analysis: keywords.extend(["soluble", "hydrophilic"])
+        if "소수성" in analysis: keywords.extend(["hydrophobic"])
+        
+        # 구조 키워드
+        if "알파" in analysis or "나선" in analysis: keywords.append("helix")
+        if "베타" in analysis or "시트" in analysis: keywords.append("sheet")
+        if "루프" in analysis: keywords.append("loop")
+        
+        return list(set(keywords))  # 중복 제거
+    except Exception as e:
+        print(f"키워드 추출 중 오류: {str(e)}")
+        return []
+
+def calculate_similarity(keywords, entry):
+    """키워드와 데이터셋 항목 간의 유사도 계산"""
+    try:
+        score = 0
+        # 시퀀스 특성 검사
+        sequence = entry['sequence'].lower()
+        for keyword in keywords:
+            if keyword in entry.get('description', '').lower():
+                score += 2
+            if keyword in sequence:
+                score += 1
+            if 'secondary_structure' in entry:
+                if keyword in ['helix'] and 'H' in entry['secondary_structure']:
+                    score += 1
+                if keyword in ['sheet'] and 'E' in entry['secondary_structure']:
+                    score += 1
+                if keyword in ['loop'] and 'L' in entry['secondary_structure']:
+                    score += 1
+        return score
+    except Exception as e:
+        print(f"유사도 계산 중 오류: {str(e)}")
+        return 0
+
+def extract_parameters_from_chat(chat_response):
+    """챗봇 응답에서 파라미터 추출"""
+    try:
+        params = {
+            'sequence_length': 100,
+            'helix_bias': 0.02,
+            'strand_bias': 0.02,
+            'loop_bias': 0.1,
+            'hydrophobic_target_score': 0
+        }
+        
+        # 응답 텍스트에서 값 추출
+        if "길이:" in chat_response:
+            length_match = re.search(r'길이: (\d+)', chat_response)
+            if length_match:
+                params['sequence_length'] = int(length_match.group(1))
+                
+        if "알파 헬릭스 비율:" in chat_response:
+            helix_match = re.search(r'알파 헬릭스 비율: ([\d.]+)', chat_response)
+            if helix_match:
+                params['helix_bias'] = float(helix_match.group(1)) / 100
+                
+        # 나머지 파라미터도 유사하게 추출
+        
+        return params
+    except Exception as e:
+        print(f"파라미터 추출 중 오류: {str(e)}")
+        return None
+
 with gr.Blocks(theme='ParityError/Interstellar') as demo:
     with gr.Row():
         with gr.Column(scale=1):
@@ -1172,7 +1326,46 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
             ]
         )
 
+    # 이벤트 핸들러 연결
+    msg.submit(
+        fn=process_chat_and_generate,
+        inputs=[msg, chatbot],
+        outputs=[
+            chatbot,
+            hero_stats,
+            hero_description,
+            output_seq,
+            output_pdb,
+            output_viewer,
+            plddt_plot
+        ]
+    )
 
+    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
+        ]
+    )
+
+    # 채팅 내용 지우기
+    clear.click(
+        lambda: (None, None, None, None, None, None, None),
+        None,
+        [chatbot, hero_stats, hero_description, output_seq, output_pdb, output_viewer, plddt_plot],
+        queue=False
+    )
+
+
+
+    
     # 챗봇 응답에 따른 결과 업데이트
     msg.submit(
         update_protein_display,