Update app.py

app.py

@@ -1,8 +1,16 @@
 import os,sys
 from huggingface_hub import HfApi
+from dotenv import load_dotenv
 
-# Hugging Face 토큰 설정
-HF_TOKEN = "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"  # 실제 토큰으로 교체
+# 환경 변수 로드 및 토큰 설정
+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"
 
 # install required packages
 os.system('pip install -q transformers')
@@ -13,7 +21,7 @@ os.environ["DGLBACKEND"] = "pytorch"
 print('Modules installed')
 
 # 필수 라이브러리 임포트
-from transformers import pipeline  # transformers 임포트 추가
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from datasets import load_dataset
 import plotly.graph_objects as go
 import numpy as np
@@ -28,10 +36,16 @@ from utils.parsers_inference import parse_pdb
 from model.util import writepdb
 from utils.inpainting_util import *
 
-# LLM과 데이터셋 로드
+# 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="CohereForAI/c4ai-command-r7b-12-2024",
-               token=HF_TOKEN)
+               model=model,
+               tokenizer=tokenizer,
+               trust_remote_code=True)
+
+# 데이터셋 로드
 ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
                  token=HF_TOKEN)