Update app.py

app.py

@@ -239,20 +239,58 @@ def dna_embed(seq: str, model_name: str = "zhihan1996/DNABERT-2-117M") -> Dict:
         try:
             import einops
         except ImportError:
-            return {"error": "einops package required. Installing..."}
+            return {"error": "einops package required. Please wait for installation and refresh the page."}
         
-        # DNABERT-2 특별 처리
-        from transformers import BertModel, BertTokenizer
+        # 간단한 대안: 더 안정적인 모델 사용
+        # DNABERT-2가 문제를 일으키면 기본 BERT 사용
+        try:
+            from transformers import AutoTokenizer, AutoModel
+            tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+            model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+        except Exception as model_error:
+            # 대체 모델 사용
+            try:
+                from transformers import BertTokenizer, BertModel
+                # 기본 BERT 모델로 폴백
+                fallback_model = "bert-base-uncased"
+                tokenizer = BertTokenizer.from_pretrained(fallback_model)
+                model = BertModel.from_pretrained(fallback_model)
+                st.warning(f"DNABERT-2 로딩 실패. 대체 모델 사용중: {fallback_model}")
+            except:
+                return {"error": f"모델 로딩 실패: {str(model_error)}"}
         
-        tokenizer = BertTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        model = BertModel.from_pretrained(model_name, trust_remote_code=True)
         model.eval()
         
+        # DNA 서열을 k-mer로 변환 (DNABERT 스타일)
+        def seq_to_kmer(seq, k=6):
+            """DNA 서열을 k-mer로 변환"""
+            kmers = []
+            for i in range(len(seq) - k + 1):
+                kmers.append(seq[i:i+k])
+            return ' '.join(kmers)
+        
+        # k-mer 변환 또는 직접 사용
+        if len(seq) > 6:
+            input_seq = seq_to_kmer(seq, k=6)
+        else:
+            input_seq = seq
+        
         with torch.no_grad():
-            inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=512)
+            inputs = tokenizer(
+                input_seq, 
+                return_tensors="pt", 
+                truncation=True, 
+                max_length=512,
+                padding=True
+            )
             outputs = model(**inputs)
-            hidden = outputs.last_hidden_state.mean(dim=1).squeeze(0)
-            vec = hidden.cpu().numpy()
+            
+            # last_hidden_state 또는 pooler_output 사용
+            if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None:
+                vec = outputs.pooler_output.squeeze(0).cpu().numpy()
+            else:
+                hidden = outputs.last_hidden_state.mean(dim=1).squeeze(0)
+                vec = hidden.cpu().numpy()
         
         # 메모리 정리
         del model
@@ -264,8 +302,9 @@ def dna_embed(seq: str, model_name: str = "zhihan1996/DNABERT-2-117M") -> Dict:
             "embedding": vec.tolist()[:10],  # 미리보기용 첫 10개만
             "size": vec.shape[0]
         }
+        
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": f"분석 중 오류 발생: {str(e)[:200]}"}
 
 def build_context(query: str, docs: List[str], index, model, use_web: bool, web_k: int) -> Tuple[str, List[Dict]]:
     """Build context from sources"""
@@ -363,7 +402,8 @@ with st.sidebar:
     )
     dna_model = st.text_input(
         "DNA Model", 
-        value="zhihan1996/DNABERT-2-117M"
+        value="bert-base-uncased",  # 더 안정적인 기본 모델
+        help="Options: bert-base-uncased (stable), zhihan1996/DNABERT-2-117M (specialized but may require more memory)"
     )
     
     use_web = st.checkbox("Enable web search", value=True)