Update app.py

app.py

@@ -1,29 +1,15 @@
 import os,sys
+from openai import OpenAI
+import gradio as gr
 
 # install required packages
-os.system('pip install -q python-dotenv')  # python-dotenv 설치 추가
-os.system('pip install -q transformers')
 os.system('pip install -q plotly')
 os.system('pip install -q 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')
 
-from huggingface_hub import HfApi
-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. Please set it in .env file")
-
-# Hugging Face API 설정
-os.environ["HUGGINGFACE_TOKEN"] = HF_TOKEN
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
-
 # 필수 라이브러리 임포트
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from datasets import load_dataset
 import plotly.graph_objects as go
 import numpy as np
@@ -38,18 +24,52 @@ from utils.parsers_inference import parse_pdb
 from model.util import writepdb
 from utils.inpainting_util import *
 
-# Cohere 모델 사용 (토큰 인증 포함)
-model_name = "CohereForAI/c4ai-command-r-plus-08-2024"
-tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, token=HF_TOKEN, trust_remote_code=True)
-pipe = pipeline("text-generation", 
-               model=model,
-               tokenizer=tokenizer,
-               trust_remote_code=True)
+# Hugging Face 토큰 설정
+ACCESS_TOKEN = os.getenv("HF_TOKEN")
+if not ACCESS_TOKEN:
+    raise ValueError("HF_TOKEN not found in environment variables")
+
+# OpenAI 클라이언트 설정 (Hugging Face 엔드포인트 사용)
+client = OpenAI(
+    base_url="https://api-inference.huggingface.co/v1/",
+    api_key=ACCESS_TOKEN,
+)
 
 # 데이터셋 로드
 ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
-                 token=HF_TOKEN)
+                 token=ACCESS_TOKEN)
+
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [{"role": "system", "content": system_message}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    
+    for message in client.chat.completions.create(
+        model="CohereForAI/c4ai-command-r-plus-08-2024",
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+        messages=messages,
+    ):
+        token = message.choices[0].delta.content
+        response += token
+        yield response
 
 # 챗봇 및 단백질 생성 관련 함수들
 def process_chat(message, history):
@@ -563,19 +583,41 @@ 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(scale=1):
             # 챗봇 인터페이스
             gr.Markdown("# 🤖 AI 단백질 설계 도우미")
-            chatbot = gr.Chatbot()
-            msg = gr.Textbox(
-                label="명령을 입력하세요 (예: COVID-19를 치료할 수 있는 단백질을 생성해주세요)",
-                placeholder="여기에 입력하세요..."
-            )
-            clear = gr.Button("대화 내용 지우기")
+            chatbot = gr.Chatbot(height=600)
+            
+            with gr.Accordion("채팅 설정", open=False):
+                system_message = gr.Textbox(
+                    value="당신은 단백질 설계를 도와주는 전문가입니다.",
+                    label="시스템 메시지"
+                )
+                max_tokens = gr.Slider(
+                    minimum=1,
+                    maximum=2048,
+                    value=512,
+                    step=1,
+                    label="최대 토큰 수"
+                )
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=4.0,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.95,
+                    step=0.05,
+                    label="Top-P"
+                )
+
 
             # 탭 인터페이스
             with gr.Tabs():
@@ -907,6 +949,17 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
         outputs=[hero_stats, hero_description, output_seq, output_pdb, output_viewer, plddt_plot]
     )
 
+    chat_interface = gr.ChatInterface(
+        respond,
+        additional_inputs=[
+            system_message,
+            max_tokens,
+            temperature,
+            top_p,
+        ],
+        chatbot=chatbot,
+    )
+
 # 실행
 demo.queue()
 demo.launch(debug=True)