Update app.py

app.py

@@ -16,7 +16,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from mergekit.config import MergeConfiguration
 from mergekit.merge import Mergekit
 from spectrum import SpectrumAnalyzer
-import distilkit
 import yaml
 from dataclasses import dataclass
 from typing import Optional, List
@@ -78,8 +77,7 @@ project_seeds = {
 - Streamlit 🌐
 - Torch 🔥
 - Transformers 🤖
-2. DistillKit MergeKit Spectrum
-- DistillKit 🧪
+2. MergeKit Spectrum
 - MergeKit 🔄
 - Spectrum 📊
 3. Transformers Diffusers Datasets
@@ -133,18 +131,6 @@ class ModelBuilder:
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         return self
 
-    @pipeline_stage
-    def apply_distillation(self, teacher_model: str, output_dir: str):
-        """Apply DistilKit for model distillation"""
-        distiller = distilkit.Distiller(
-            teacher_model=teacher_model,
-            student_model=self.model,
-            tokenizer=self.tokenizer
-        )
-        distiller.distill(output_dir=output_dir)
-        self.model = distiller.student_model
-        return self
-
     @pipeline_stage
     def apply_merge(self, models_to_merge: List[str], output_dir: str):
         """Apply Mergekit for model merging"""
@@ -416,7 +402,6 @@ if st.button("Predict"):
             )
             model_name = st.text_input("Model Name", "custom-model")
             domain = st.text_input("Target Domain", "general")
-            use_distillation = st.checkbox("Apply Distillation", True)
             use_merging = st.checkbox("Apply Model Merging", False)
             use_spectrum = st.checkbox("Apply Spectrum Specialization", True)
 
@@ -433,13 +418,6 @@ if st.button("Predict"):
             with st.status("Building advanced model...", expanded=True) as status:
                 builder.load_base_model(config.base_model)
                 
-                if use_distillation:
-                    teacher_model = st.selectbox(
-                        "Select Teacher Model",
-                        ["mistral-13b", "llama-2-13b"]
-                    )
-                    builder.apply_distillation(teacher_model, f"distilled_{config.name}")
-                
                 if use_merging:
                     models_to_merge = st.multiselect(
                         "Select Models to Merge",
@@ -471,7 +449,7 @@ if st.button("Generate"):
 """
             with open("advanced_app.py", "w") as f:
                 f.write(app_code)
-            reqs = "streamlit\ntorch\ntransformers\n"
+            reqs = "streamlit\ntorch\ntransformers\nmergekit\nspectrum\n"
             with open("advanced_requirements.txt", "w") as f:
                 f.write(reqs)
             readme = f"""