Update app.py

app.py

@@ -1,15 +1,25 @@
 import os,sys
+from transformers import pipeline
+from datasets import load_dataset
+import gradio as gr
+from dotenv import load_dotenv
+
+# 환경 변수 로드
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN not found in environment variables")
 
 # install required packages
-os.system('pip install plotly')  # plotly 설치
-os.system('pip install matplotlib')  # matplotlib 설치
+os.system('pip install plotly')
+os.system('pip install matplotlib')
 os.system('pip install dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html')
 os.environ["DGLBACKEND"] = "pytorch"
 print('Modules installed')
 
+# 기존 import문들 유지
 import plotly.graph_objects as go
 import numpy as np
-import gradio as gr
 import py3Dmol
 from io import StringIO
 import json
@@ -21,59 +31,56 @@ from utils.parsers_inference import parse_pdb
 from model.util import writepdb
 from utils.inpainting_util import *
 
-# install environment goods
-#os.system("pip -q install dgl -f https://data.dgl.ai/wheels/cu113/repo.html")
-os.system('pip install dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html')
-#os.system('pip install gradio')
-os.environ["DGLBACKEND"] = "pytorch"
-#os.system(f'pip install -r ./PROTEIN_GENERATOR/requirements.txt')
-print('Modules installed')
-
-#os.system('pip install --force gradio==3.36.1')
-#os.system('pip install gradio_client==0.2.7')
-#os.system('pip install \"numpy<2\"')
-#os.system('pip install numpy --upgrade')
-#os.system('pip install --force numpy==1.24.1')
-
-
-if not os.path.exists('./SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'):
-    print('Downloading model weights 1')
-    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt')
-    print('Successfully Downloaded')
-
-if not os.path.exists('./SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'):
-    print('Downloading model weights 2')
-    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt')
-    print('Successfully Downloaded')
-
-import numpy as np
-import gradio as gr
-import py3Dmol
-from io import StringIO
-import json
-import secrets
-import copy
-import matplotlib.pyplot as plt
-from utils.sampler import HuggingFace_sampler
-from utils.parsers_inference import parse_pdb
-from model.util import writepdb
-from utils.inpainting_util import *
-
-
-plt.rcParams.update({'font.size': 13})
-
-with open('./tmp/args.json','r') as f:
-    args = json.load(f)
+# 새로운 모델과 데이터셋 로드
+pipe = pipeline("text-generation", 
+               model="CohereForAI/c4ai-command-r-plus-08-2024",
+               token=HF_TOKEN)
+ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
+                 token=HF_TOKEN)
+
+# 챗봇 및 단백질 생성 관련 함수들
+def process_chat(message, history):
+    messages = [{"role": "user", "content": message}]
+    response = pipe(messages)[0]['generated_text']
+    
+    if any(keyword in message.lower() for keyword in ['protein', 'generate', '단백질', '생성']):
+        relevant_data = search_protein_data(message)
+        params = extract_parameters(response, relevant_data)
+        protein_result = generate_protein(params)
+        explanation = generate_explanation(protein_result, params)
+        return response + "\n\n" + explanation
+    
+    return response
+
+def search_protein_data(query):
+    relevant_entries = []
+    for entry in ds['train']:
+        if any(keyword in entry['sequence'].lower() for keyword in query.lower().split()):
+            relevant_entries.append(entry)
+    return relevant_entries
+
+def extract_parameters(llm_response, dataset_info):
+    params = {
+        'sequence_length': 100,
+        'helix_bias': 0.02,
+        'strand_bias': 0.02,
+        'loop_bias': 0.1,
+        'hydrophobic_target_score': 0
+    }
+    return params
+
+def generate_explanation(result, params):
+    explanation = f"""
+    생성된 단백질 분석:
+    - 길이: {params['sequence_length']} 아미노산
+    - 구조적 특징:
+        * 알파 나선 비율: {params['helix_bias']*100}%
+        * 베타 시트 비율: {params['strand_bias']*100}%
+        * 루프 구조 비율: {params['loop_bias']*100}%
+    - 특수 기능: {result.get('special_features', '없음')}
+    """
+    return explanation
 
-# manually set checkpoint to load
-args['checkpoint'] = None
-args['dump_trb'] = False
-args['dump_args'] = True
-args['save_best_plddt'] = True
-args['T'] = 25
-args['strand_bias'] = 0.0
-args['loop_bias'] = 0.0
-args['helix_bias'] = 0.0
 
 def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias, 
                     secondary_structure, aa_bias, aa_bias_potential, 
@@ -543,11 +550,21 @@ def combined_generation(name, strength, flexibility, speed, defense, size, abili
             None
         )
 
+# Gradio 인터페이스 수정
 with gr.Blocks(theme='ParityError/Interstellar') as demo:
     with gr.Row():
-        with gr.Column():
-            gr.Markdown("# 🦸‍♂️ 슈퍼히어로 단백질 만들기")
-            
+        # 왼쪽 열: 챗봇 및 컨트롤 패널
+        with gr.Column(scale=1):
+            # 챗봇 인터페이스
+            gr.Markdown("# 🤖 AI 단백질 설계 도우미")
+            chatbot = gr.Chatbot()
+            msg = gr.Textbox(
+                label="명령을 입력하세요 (예: COVID-19를 치료할 수 있는 단백질을 생성해주세요)",
+                placeholder="여기에 입력하세요..."
+            )
+            clear = gr.Button("대화 내용 지우기")
+
+            # 탭 인터페이스
             with gr.Tabs():
                 with gr.TabItem("🦸‍♂️ 히어로 디자인"):
                     gr.Markdown("""
@@ -616,9 +633,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                     # 생성 버튼
                     create_btn = gr.Button("🧬 히어로 생성!", variant="primary", scale=2)
 
-
-
-
                 with gr.TabItem("🧬 히어로 DNA 설계"):
                     gr.Markdown("""
                     ### 🧪 히어로 DNA 고급 설정
@@ -681,7 +695,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                                 visible=True
                             )
                     
-                    # 아미노산 구성 설정 추가
                     with gr.Accordion(label='🧬 DNA 구성 설정', open=False):
                         gr.Markdown("""
                         특정 아미노산의 비율을 조절하여 히어로의 특성을 강화할 수 있습니다.
@@ -699,7 +712,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                                 placeholder='1.0-5.0 사이 값 입력'
                             )
 
-                    # 환경 적응력 설정 추가
                     with gr.Accordion(label='🌍 환경 적응력 설정', open=False):
                         gr.Markdown("""
                         히어로의 환경 적응력을 조절합니다.
@@ -717,7 +729,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                                 placeholder='1.0-2.0 사이 값 입력'
                             )
 
-                    # 확산 매개변수 설정
                     with gr.Accordion(label='⚙️ 고급 설정', open=False):
                         gr.Markdown("""
                         DNA 생성 과정의 세부 매개변수를 조정합니다.
@@ -734,9 +745,7 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                                 value='normal'
                             )
                     
-                    design_btn = gr.Button("🧬 DNA 설계 생성!", variant="primary", scale=2)                
-
-
+                    design_btn = gr.Button("🧬 DNA 설계 생성!", variant="primary", scale=2)
 
                 with gr.TabItem("🧪 히어로 유전자 강화"):
                     gr.Markdown("""
@@ -792,21 +801,14 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                     pssm = gr.File(label='분석된 가문 특성')
                     pssm_view = gr.Plot(label='가문 특성 분석 결과')
                     pssm_gen_btn = gr.Button("✨ 가문 특성 분석", variant="secondary")
-                    inherit_btn = gr.Button("👑 가문의 힘 계승!", variant="primary", scale=2)             
-
+                    inherit_btn = gr.Button("👑 가문의 힘 계승!", variant="primary", scale=2)
 
-
-
-
-        with gr.Column():
+        # 오른쪽 열: 결과 표시
+        with gr.Column(scale=1):
             gr.Markdown("## 🦸‍♂️ 히어로 프로필")
-            
-            # 능력치 레이더 차트
             hero_stats = gr.Plot(label="능력치 분석")
-            
-            # 히어로 설명
             hero_description = gr.Textbox(label="히어로 특성", lines=3)
-
+            
             gr.Markdown("## 🧬 히어로 DNA 분석 결과")
             gr.Markdown("#### ⚡ DNA 안정성 점수")
             plddt_plot = gr.Plot(label='안정성 분석')
@@ -818,6 +820,11 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
             output_viewer = gr.HTML()
 
     # 이벤트 연결
+    # 챗봇 이벤트
+    msg.submit(process_chat, [msg, chatbot], [chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+    # UI 컨트롤 이벤트
     seq_opt.change(
         fn=toggle_seq_input,
         inputs=[seq_opt],
@@ -832,85 +839,61 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
         queue=False
     )
 
-    preview_btn.click(get_motif_preview,[pdb_id_code, contigs],[preview_viewer, rewrite_pdb])
-    pssm_gen_btn.click(get_pssm,[fasta_msa,input_pssm],[pssm_view, pssm])
-
-    # 각 탭의 생성 버튼 연결
-    create_btn.click(
-        combined_generation,
-        inputs=[
-            hero_name, strength, flexibility, speed, defense, hero_size, special_ability,
-            sequence, seq_len, helix_bias, strand_bias, loop_bias, 
-            secondary_structure, aa_bias, aa_bias_potential, 
-            num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
-            contigs, pssm, seq_mask, str_mask, rewrite_pdb
-        ],
-        outputs=[
-            hero_stats, 
-            hero_description, 
-            output_seq,
-            output_pdb,
-            output_viewer,
-            plddt_plot
-        ]
+    preview_btn.click(
+        get_motif_preview,
+        inputs=[pdb_id_code, contigs],
+        outputs=[preview_viewer, rewrite_pdb]
     )
 
-    design_btn.click(
-        combined_generation,
-        inputs=[
-            hero_name, strength, flexibility, speed, defense, hero_size, special_ability,
-            sequence, seq_len, helix_bias, strand_bias, loop_bias, 
-            secondary_structure, aa_bias, aa_bias_potential, 
-            num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
-            contigs, pssm, seq_mask, str_mask, rewrite_pdb
-        ],
-        outputs=[
-            hero_stats, 
-            hero_description, 
-            output_seq,
-            output_pdb,
-            output_viewer,
-            plddt_plot
-        ]
+    pssm_gen_btn.click(
+        get_pssm,
+        inputs=[fasta_msa, input_pssm],
+        outputs=[pssm_view, pssm]
     )
 
-    enhance_btn.click(
-        combined_generation,
-        inputs=[
-            hero_name, strength, flexibility, speed, defense, hero_size, special_ability,
-            sequence, seq_len, helix_bias, strand_bias, loop_bias, 
-            secondary_structure, aa_bias, aa_bias_potential, 
-            num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
-            contigs, pssm, seq_mask, str_mask, rewrite_pdb
-        ],
-        outputs=[
-            hero_stats, 
-            hero_description, 
-            output_seq,
-            output_pdb,
-            output_viewer,
-            plddt_plot
-        ]
-    )
+    # 챗봇 기반 단백질 생성 결과 업데이트
+    def update_protein_display(chat_response):
+        if "생성된 단백질 분석" in chat_response:
+            params = extract_parameters_from_chat(chat_response)
+            result = generate_protein(params)
+            return {
+                hero_stats: create_radar_chart(calculate_hero_stats(params)),
+                hero_description: chat_response,
+                output_seq: result[0],
+                output_pdb: result[1],
+                output_viewer: display_pdb(result[1]),
+                plddt_plot: result[3]
+            }
+        return None
+
+    # 각 생성 버튼 이벤트 연결
+    for btn in [create_btn, design_btn, enhance_btn, inherit_btn]:
+        btn.click(
+            combined_generation,
+            inputs=[
+                hero_name, strength, flexibility, speed, defense, hero_size, special_ability,
+                sequence, seq_len, helix_bias, strand_bias, loop_bias, 
+                secondary_structure, aa_bias, aa_bias_potential, 
+                num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
+                contigs, pssm, seq_mask, str_mask, rewrite_pdb
+            ],
+            outputs=[
+                hero_stats, 
+                hero_description, 
+                output_seq,
+                output_pdb,
+                output_viewer,
+                plddt_plot
+            ]
+        )
 
-    inherit_btn.click(
-        combined_generation,
-        inputs=[
-            hero_name, strength, flexibility, speed, defense, hero_size, special_ability,
-            sequence, seq_len, helix_bias, strand_bias, loop_bias, 
-            secondary_structure, aa_bias, aa_bias_potential, 
-            num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
-            contigs, pssm, seq_mask, str_mask, rewrite_pdb
-        ],
-        outputs=[
-            hero_stats, 
-            hero_description, 
-            output_seq,
-            output_pdb,
-            output_viewer,
-            plddt_plot
-        ]
+    # 챗봇 응답에 따른 결과 업데이트
+    msg.submit(
+        update_protein_display,
+        inputs=[chatbot],
+        outputs=[hero_stats, hero_description, output_seq, output_pdb, output_viewer, plddt_plot]
     )
 
+# 실행
 demo.queue()
-demo.launch(debug=True)                    
+demo.launch(debug=True)