Update app.py

app.py

@@ -211,13 +211,19 @@ def esm2_embed(seq: str, model_name: str = "facebook/esm2_t6_8M_UR50D") -> Dict:
         model.eval()
         
         with torch.no_grad():
-            inputs = tokenizer(seq, return_tensors="pt")
+            inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=1024)
             outputs = model(**inputs, output_hidden_states=True)
             hidden = outputs.hidden_states[-1].mean(dim=1).squeeze(0)
-            vec = hidden.numpy()
+            vec = hidden.cpu().numpy()
+        
+        # 메모리 정리
+        del model
+        del tokenizer
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
         return {
-            "embedding": vec.tolist(),
+            "embedding": vec.tolist()[:10],  # 미리보기용 첫 10개만
             "size": vec.shape[0]
         }
     except Exception as e:
@@ -229,6 +235,12 @@ def dna_embed(seq: str, model_name: str = "zhihan1996/DNABERT-2-117M") -> Dict:
         return {"error": "PyTorch/Transformers not available"}
     
     try:
+        # einops 체크
+        try:
+            import einops
+        except ImportError:
+            return {"error": "einops package required. Installing..."}
+        
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
         model.eval()
@@ -237,10 +249,16 @@ def dna_embed(seq: str, model_name: str = "zhihan1996/DNABERT-2-117M") -> Dict:
             inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=512)
             outputs = model(**inputs)
             hidden = outputs.last_hidden_state.mean(dim=1).squeeze(0)
-            vec = hidden.numpy()
+            vec = hidden.cpu().numpy()
+        
+        # 메모리 정리
+        del model
+        del tokenizer
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
         return {
-            "embedding": vec.tolist(),
+            "embedding": vec.tolist()[:10],  # 미리보기용 첫 10개만
             "size": vec.shape[0]
         }
     except Exception as e: