Update app.py

app.py

@@ -54,18 +54,18 @@ class ExperimentDatabase:
     def __init__(self, db_path: str):
         self.db_path = db_path
         self.init_database()
+        self.migrate_database()  # 마이그레이션 실행
     
     def init_database(self):
         """데이터베이스 초기화"""
         with sqlite3.connect(self.db_path) as conn:
             cursor = conn.cursor()
             
-            # 실험 테이블
+            # 실험 테이블 (기본 버전)
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS experiments (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
                     model_type TEXT NOT NULL,
-                    base_model_url TEXT,
                     sequence_length INTEGER,
                     power_mode TEXT,
                     compression_level REAL,
@@ -92,14 +92,40 @@ class ExperimentDatabase:
                 ON experiments(timestamp DESC)
             """)
             
-            cursor.execute("""
-                CREATE INDEX IF NOT EXISTS idx_base_model 
-                ON experiments(base_model_url)
-            """)
-            
             conn.commit()
             print("✅ Database initialized")
     
+    def migrate_database(self):
+        """데이터베이스 마이그레이션 - 새 컬럼 추가"""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            
+            # 컬럼 존재 확인
+            cursor.execute("PRAGMA table_info(experiments)")
+            columns = [column[1] for column in cursor.fetchall()]
+            
+            # base_model_url 컬럼이 없으면 추가
+            if 'base_model_url' not in columns:
+                try:
+                    cursor.execute("""
+                        ALTER TABLE experiments 
+                        ADD COLUMN base_model_url TEXT
+                    """)
+                    print("✅ Database migrated: base_model_url column added")
+                except sqlite3.OperationalError as e:
+                    print(f"⚠️ Migration warning: {e}")
+            
+            # 인덱스 추가
+            try:
+                cursor.execute("""
+                    CREATE INDEX IF NOT EXISTS idx_base_model 
+                    ON experiments(base_model_url)
+                """)
+            except sqlite3.OperationalError:
+                pass
+            
+            conn.commit()
+    
     def save_experiment(self, config: Dict, metrics: Dict) -> int:
         """실험 저장"""
         with sqlite3.connect(self.db_path) as conn:
@@ -171,13 +197,17 @@ class ExperimentDatabase:
             """)
             by_model = dict(cursor.fetchall())
             
-            cursor.execute("""
-                SELECT base_model_url, COUNT(*) as count 
-                FROM experiments 
-                WHERE base_model_url IS NOT NULL
-                GROUP BY base_model_url
-            """)
-            by_base_model = dict(cursor.fetchall())
+            # base_model_url 컬럼이 있는 경우에만 조회
+            try:
+                cursor.execute("""
+                    SELECT base_model_url, COUNT(*) as count 
+                    FROM experiments 
+                    WHERE base_model_url IS NOT NULL
+                    GROUP BY base_model_url
+                """)
+                by_base_model = dict(cursor.fetchall())
+            except sqlite3.OperationalError:
+                by_base_model = {}
             
             return {
                 'total_experiments': total,
@@ -189,48 +219,66 @@ class RetentionVectorStore:
     """ChromaDB 벡터 저장소"""
     
     def __init__(self, persist_directory: str):
-        self.client = chromadb.Client(Settings(
-            persist_directory=persist_directory,
-            anonymized_telemetry=False
-        ))
-        
-        self.collection = self.client.get_or_create_collection(
-            name="retention_states",
-            metadata={"description": "PHOENIX Retention states"}
-        )
-        print("✅ Vector store initialized")
+        try:
+            self.client = chromadb.Client(Settings(
+                persist_directory=persist_directory,
+                anonymized_telemetry=False
+            ))
+            
+            self.collection = self.client.get_or_create_collection(
+                name="retention_states",
+                metadata={"description": "PHOENIX Retention states"}
+            )
+            print("✅ Vector store initialized")
+        except Exception as e:
+            print(f"⚠️ Vector store initialization warning: {e}")
+            self.client = None
+            self.collection = None
     
     def add_retention_state(self, experiment_id: int, states: Dict, metadata: Dict):
         """Retention state 저장"""
-        state_vector = self._states_to_vector(states)
+        if self.collection is None:
+            return
         
-        self.collection.add(
-            embeddings=[state_vector.tolist()],
-            metadatas=[{**metadata, 'experiment_id': experiment_id}],
-            ids=[f"exp_{experiment_id}"]
-        )
+        try:
+            state_vector = self._states_to_vector(states)
+            
+            self.collection.add(
+                embeddings=[state_vector.tolist()],
+                metadatas=[{**metadata, 'experiment_id': experiment_id}],
+                ids=[f"exp_{experiment_id}"]
+            )
+        except Exception as e:
+            print(f"⚠️ Vector store save warning: {e}")
     
     def search(self, query: str, top_k: int = 10) -> List[Dict]:
         """실험 검색"""
-        query_vector = self._text_to_vector(query)
-        
-        results = self.collection.query(
-            query_embeddings=[query_vector.tolist()],
-            n_results=top_k
-        )
-        
-        if not results['ids'][0]:
+        if self.collection is None:
             return []
         
-        formatted_results = []
-        for i in range(len(results['ids'][0])):
-            formatted_results.append({
-                'experiment_id': results['metadatas'][0][i].get('experiment_id'),
-                'score': 1.0 - results['distances'][0][i],
-                'metadata': results['metadatas'][0][i]
-            })
-        
-        return formatted_results
+        try:
+            query_vector = self._text_to_vector(query)
+            
+            results = self.collection.query(
+                query_embeddings=[query_vector.tolist()],
+                n_results=top_k
+            )
+            
+            if not results['ids'][0]:
+                return []
+            
+            formatted_results = []
+            for i in range(len(results['ids'][0])):
+                formatted_results.append({
+                    'experiment_id': results['metadatas'][0][i].get('experiment_id'),
+                    'score': 1.0 - results['distances'][0][i],
+                    'metadata': results['metadatas'][0][i]
+                })
+            
+            return formatted_results
+        except Exception as e:
+            print(f"⚠️ Vector store search warning: {e}")
+            return []
     
     def _states_to_vector(self, states: Dict) -> np.ndarray:
         """States를 고정 크기 벡터로 변환"""
@@ -648,14 +696,18 @@ def initialize_default_models():
     models = {}
     
     try:
-        # PHOENIX with Granite
-        models['phoenix_granite'] = PHOENIXRetention(
-            d_model=512,
-            d_state=256,
-            num_layers=12,
-            device=DEVICE,
-            base_model_url=DEFAULT_MODEL
-        )
+        # PHOENIX with Granite (옵션)
+        try:
+            models['phoenix_granite'] = PHOENIXRetention(
+                d_model=512,
+                d_state=256,
+                num_layers=12,
+                device=DEVICE,
+                base_model_url=DEFAULT_MODEL
+            )
+            print("✅ phoenix_granite initialized")
+        except Exception as e:
+            print(f"⚠️ phoenix_granite initialization skipped: {e}")
         
         # PHOENIX without base
         models['phoenix_standalone'] = PHOENIXRetention(
@@ -665,21 +717,32 @@ def initialize_default_models():
             device=DEVICE,
             base_model_url=None
         )
+        print("✅ phoenix_standalone initialized")
         
-        # Transformer baseline
-        models['transformer_granite'] = TransformerBaseline(
-            d_model=512,
-            d_state=256,
-            device=DEVICE,
-            base_model_url=DEFAULT_MODEL
-        )
+        # Transformer baseline (옵션)
+        try:
+            models['transformer_granite'] = TransformerBaseline(
+                d_model=512,
+                d_state=256,
+                device=DEVICE,
+                base_model_url=DEFAULT_MODEL
+            )
+            print("✅ transformer_granite initialized")
+        except Exception as e:
+            print(f"⚠️ transformer_granite initialization skipped: {e}")
         
-        print("✅ Default models initialized")
+        print(f"✅ {len(models)} models initialized successfully")
         return models
         
     except Exception as e:
         print(f"❌ Model initialization failed: {e}")
-        return {}
+        return {'phoenix_standalone': PHOENIXRetention(
+            d_model=512,
+            d_state=256,
+            num_layers=12,
+            device=DEVICE,
+            base_model_url=None
+        )}
 
 # 데이터베이스 및 모델 초기화
 db = ExperimentDatabase(DB_PATH)
@@ -713,7 +776,7 @@ def run_retention_experiment(
         # 실험 설정
         config = {
             'model_type': model_name,
-            'base_model_url': custom_model_url if custom_model_url else model.base_model_url,
+            'base_model_url': custom_model_url if custom_model_url else (model.base_model_url if hasattr(model, 'base_model_url') else None),
             'sequence_length': sequence_length,
             'power_mode': power_mode,
             'compression_level': compression_level,
@@ -890,10 +953,16 @@ def view_experiment_history(limit=20):
             title='모델별 실행 시간 추이'
         )
         
+        # base_model_url 컬럼이 있는지 확인
+        if 'base_model_url' in df.columns:
+            display_cols = ['id', 'model_type', 'base_model_url', 'sequence_length', 'elapsed_time', 'throughput', 'timestamp']
+        else:
+            display_cols = ['id', 'model_type', 'sequence_length', 'elapsed_time', 'throughput', 'timestamp']
+        
         history_text = f"""
 ## 📊 실험 이력 ({len(df)}개)
 
-{df[['id', 'model_type', 'base_model_url', 'sequence_length', 'elapsed_time', 'throughput', 'timestamp']].to_markdown(index=False)}
+{df[display_cols].to_markdown(index=False)}
         """
         
         return history_text, fig
@@ -917,9 +986,10 @@ def get_database_statistics():
         for model, count in stats['by_model'].items():
             stats_text += f"- **{model}**: {count}개\n"
         
-        stats_text += "\n### Base Model별 실험 수\n"
-        for base_model, count in stats['by_base_model'].items():
-            stats_text += f"- **{base_model}**: {count}개\n"
+        if stats['by_base_model']:
+            stats_text += "\n### Base Model별 실험 수\n"
+            for base_model, count in stats['by_base_model'].items():
+                stats_text += f"- **{base_model}**: {count}개\n"
         
         return stats_text
         
@@ -954,7 +1024,7 @@ with gr.Blocks(
                 with gr.Column(scale=1):
                     model_select = gr.Dropdown(
                         choices=list(MODELS.keys()),
-                        value='phoenix_granite',
+                        value=list(MODELS.keys())[0] if MODELS else None,
                         label="기본 모델 선택"
                     )