Update app.py

app.py

@@ -14,7 +14,6 @@ from trl import SFTTrainer
 from datasets import Dataset
 
 # --- КОНФИГУРАЦИЯ ---
-# Используем Maincoder-1B (он легкий и не убьет Space по памяти)
 MODEL_ID = "Maincode/Maincoder-1B"
 OUTPUT_DIR = "mandre_qlora_adapter"
 JSON_FILE_NAME = "train_data.json"
@@ -24,19 +23,16 @@ chat_model = None
 chat_tokenizer = None
 
 # ==========================================
-# ЧАСТЬ 1: ГЕНЕРАТОР ДАТАСЕТА (ИСПРАВЛЕННЫЙ)
+# ЧАСТЬ 1: ГЕНЕРАТОР ДАТАСЕТА
 # ==========================================
 
 def generate_json_dataset(files):
-    # Защита от пустого ввода
     if not files:
         return None, "❌ Ошибка: Вы не загрузили файлы."
 
     data_entries = []
     
-    # Gradio 4.x передает список объектов, берем пути корректно
     for file_item in files:
-        # Получаем путь к файлу (обрабатываем разные варианты ввода Gradio)
         if isinstance(file_item, str):
             file_path = file_item
         elif hasattr(file_item, 'name'):
@@ -49,11 +45,10 @@ def generate_json_dataset(files):
                 content = f.read()
         except Exception as e:
             print(f"Skipping file {file_path}: {e}")
-            continue # Пропускаем бинарники или ошибки чтения
+            continue 
             
         filename = os.path.basename(file_path)
         
-        # Формируем пару Инструкция-Ответ
         instruction = f"Analyze the code/text in file '{filename}' from the MandreLib project."
         text = f"### Instruction:\n{instruction}\n\n### Response:\n{content}<|endoftext|>"
         
@@ -62,37 +57,30 @@ def generate_json_dataset(files):
     if not data_entries:
         return None, "❌ Не удалось прочитать ни один текстовый файл."
 
-    # Сохраняем JSON
     try:
         with open(JSON_FILE_NAME, 'w', encoding='utf-8') as f:
             json.dump(data_entries, f, indent=4, ensure_ascii=False)
-        
-        # Возвращаем ПУТЬ К ФАЙЛУ (строку) и сообщение (строку)
-        # Важно: возвращаем абсолютный путь для надежности
         abs_path = os.path.abspath(JSON_FILE_NAME)
         return abs_path, f"✅ Готово! Обработано файлов: {len(data_entries)}. Файл {JSON_FILE_NAME} создан."
     except Exception as e:
         return None, f"❌ Ошибка записи JSON: {e}"
 
 # ==========================================
-# ЧАСТЬ 2: ОБУЧЕНИЕ (CPU STREAMING)
+# ЧАСТЬ 2: ОБУЧЕНИЕ (ИСПРАВЛЕНО)
 # ==========================================
 
 def train_mandre_ai(file_obj, epochs, lr):
-    # Проверка наличия файла
     if file_obj is None:
-        # Пытаемся найти файл, если он был создан ранее, но не передан через UI
         if os.path.exists(JSON_FILE_NAME):
             json_path = JSON_FILE_NAME
-            yield f"⚠️ Файл не передан в поле, но найден {JSON_FILE_NAME} на диске. Используем его."
+            yield f"⚠️ Файл не передан, используем {JSON_FILE_NAME} из прошлой генерации."
         else:
-            yield "❌ Ошибка: Загрузите JSON файл или сгенерируйте его во вкладке 1!"
+            yield "❌ Ошибка: Нет файла с данными!"
             return
     else:
-        # Gradio может передать объект файла или путь
         json_path = file_obj.name if hasattr(file_obj, 'name') else file_obj
 
-    yield f"🚀 Начинаем процесс... (Модель: {MODEL_ID})"
+    yield f"🚀 Старт обучения {MODEL_ID}..."
     
     try:
         # 1. Загрузка данных
@@ -100,13 +88,13 @@ def train_mandre_ai(file_obj, epochs, lr):
             data = json.load(f)
         dataset = Dataset.from_pandas(pd.DataFrame(data))
         
-        yield f"📊 Данные загружены: {len(dataset)} примеров."
+        yield f"📊 Данные: {len(dataset)} строк. Загрузка токенизатора..."
 
-        # 2. Токенизатор
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        # 2. Токенизатор (FIX: use_fast=False чтобы избежать ошибки Rust)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False, trust_remote_code=True)
         tokenizer.pad_token = tokenizer.eos_token
 
-        # 3. LoRA Config (Оптимизировано для CPU)
+        # 3. LoRA Config
         peft_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
             inference_mode=False,
@@ -116,7 +104,7 @@ def train_mandre_ai(file_obj, epochs, lr):
             target_modules=['q_proj', 'v_proj', 'k_proj', 'o_proj']
         )
 
-        # 4. Аргументы обучения
+        # 4. Аргументы
         training_args = TrainingArguments(
             output_dir=OUTPUT_DIR,
             num_train_epochs=float(epochs),
@@ -124,23 +112,23 @@ def train_mandre_ai(file_obj, epochs, lr):
             gradient_accumulation_steps=4,
             learning_rate=float(lr),
             weight_decay=0.01,
-            use_cpu=True,           # ПРИНУДИТЕЛЬНО CPU
-            no_cuda=True,           # ОТКЛЮЧИТЬ CUDA
-            fp16=False,             # CPU любит fp32
+            use_cpu=True,           
+            no_cuda=True,           
+            fp16=False,             
             logging_steps=1,
             save_total_limit=1,
             push_to_hub=False,
-            report_to="none"        # Отключаем wandb чтобы не спамил
+            report_to="none"
         )
 
-        yield "📥 Загрузка модели в память (это может занять минуту)..."
+        yield "📥 Загрузка модели (Maincoder-1B)..."
         
         # 5. Загрузка модели
-        model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
         model.gradient_checkpointing_enable() 
         model.enable_input_require_grads()
 
-        yield "🧠 Модель готова. Инициализация тренера..."
+        yield "🧠 Инициализация тренера..."
 
         # 6. Trainer
         trainer = SFTTrainer(
@@ -153,20 +141,19 @@ def train_mandre_ai(file_obj, epochs, lr):
             max_seq_length=1024
         )
 
-        yield "🔥 ОБУЧЕНИЕ ЗАПУЩЕНО! Процесс пошел. Это будет долго, не закрывайте вкладку."
+        yield "🔥 ОБУЧЕНИЕ ЗАПУЩЕНО! Ждите завершения..."
         
-        # Запуск обучения
         trainer.train()
         
-        yield "💾 Сохранение адаптера..."
+        yield "💾 Сохранение..."
         trainer.model.save_pretrained(OUTPUT_DIR)
         tokenizer.save_pretrained(OUTPUT_DIR)
         
-        yield f"✅ УРА! Обучение завершено. Адаптер сохранен в папку '{OUTPUT_DIR}'. Переходите в Чат."
+        yield f"✅ УСПЕХ! Адаптер в папке '{OUTPUT_DIR}'. Можно чатиться."
 
     except Exception as e:
         import traceback
-        yield f"❌ КРИТИЧЕСКАЯ ОШИБКА:\n{traceback.format_exc()}"
+        yield f"❌ ОШИБКА:\n{traceback.format_exc()}"
 
 # ==========================================
 # ЧАСТЬ 3: ЧАТ
@@ -177,18 +164,19 @@ def load_chat_model():
     if chat_model is not None: return "Уже загружено"
 
     try:
-        chat_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        # FIX: use_fast=False и здесь
+        chat_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False, trust_remote_code=True)
         
         if os.path.exists(os.path.join(OUTPUT_DIR, "adapter_config.json")):
-            base = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+            base = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
             chat_model = PeftModel.from_pretrained(base, OUTPUT_DIR)
-            return f"✅ Загружен адаптер QLoRA из {OUTPUT_DIR}!"
+            return f"✅ Адаптер QLoRA загружен!"
         else:
-            chat_model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
-            return "⚠️ Адаптер не найден (не обучили?). Загружена 'чистая' модель."
+            chat_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
+            return "⚠️ Адаптер не найден. Работает чистая модель."
             
     except Exception as e:
-        return f"Ошибка загрузки: {e}"
+        return f"Ошибка: {e}"
 
 def generate_answer(prompt, history):
     if not chat_model: 
@@ -198,7 +186,6 @@ def generate_answer(prompt, history):
     formatted_prompt = f"### Instruction:\n{prompt}\n\n### Response:\n"
     inputs = chat_tokenizer(formatted_prompt, return_tensors="pt")
     
-    # Генерация
     outputs = chat_model.generate(
         **inputs, 
         max_new_tokens=300,
@@ -216,52 +203,33 @@ def generate_answer(prompt, history):
 # ИНТЕРФЕЙС
 # ==========================================
 
-with gr.Blocks(title="MandreAI 1B CPU Trainer") as demo:
-    gr.Markdown("# 🦎 MandreAI (Stable CPU Version)")
-    gr.Markdown("Версия для обучения на бесплатном процессоре Hugging Face.")
+with gr.Blocks(title="MandreAI Fix") as demo:
+    gr.Markdown("# 🦎 MandreAI 1B (CPU Fix)")
     
     with gr.Tabs():
-        # Вкладка 1
         with gr.Tab("1. Датасет"):
-            gr.Markdown("Загрузите файлы `.py`, `.md`, `.java`.")
             files_input = gr.File(file_count="multiple", label="Исходные файлы")
             btn_gen = gr.Button("Создать JSON", variant="primary")
-            
-            # ВАЖНО: Определяем выходы точно так, как возвращает функция (File, Textbox)
-            json_output = gr.File(label="Готовый датасет (скачается автоматически)")
-            status_output = gr.Textbox(label="Статус генерации")
-            
-            btn_gen.click(
-                generate_json_dataset, 
-                inputs=[files_input], 
-                outputs=[json_output, status_output]
-            )
+            json_output = gr.File(label="Готовый датасет")
+            status_output = gr.Textbox(label="Статус")
+            btn_gen.click(generate_json_dataset, inputs=[files_input], outputs=[json_output, status_output])
 
-        # Вкладка 2
         with gr.Tab("2. Обучение"):
-            gr.Markdown("Загрузите полученный `.json` файл сюда (или он подхватится автоматически, если создан).")
             with gr.Row():
-                train_file_input = gr.File(label="Файл train_data.json")
+                train_file_input = gr.File(label="train_data.json")
                 epochs = gr.Number(value=3, label="Эпохи", precision=0)
-                lr = gr.Number(value=2e-4, label="Learning Rate")
+                lr = gr.Number(value=2e-4, label="LR")
             
             btn_train = gr.Button("ЗАПУСТИТЬ ОБУЧЕНИЕ", variant="stop")
-            # Используем Textbox как лог
-            log_output = gr.Textbox(label="Лог процесса (обновляется в реальном времени)", lines=10, interactive=False)
+            log_output = gr.Textbox(label="Лог", lines=10)
             
-            btn_train.click(
-                train_mandre_ai, 
-                inputs=[train_file_input, epochs, lr], 
-                outputs=[log_output]
-            )
+            btn_train.click(train_mandre_ai, inputs=[train_file_input, epochs, lr], outputs=[log_output])
 
-        # Вкладка 3
         with gr.Tab("3. Чат"):
             chatbot = gr.Chatbot(label="MandreAI")
-            msg_input = gr.Textbox(label="Ваш вопрос")
+            msg_input = gr.Textbox(label="Вопрос")
             btn_send = gr.Button("Отправить")
-            
             btn_send.click(generate_answer, [msg_input, chatbot], chatbot)
 
 if __name__ == "__main__":
-    demo.queue().launch(allowed_paths=["."]) # Разрешаем доступ к локальным файлам
+    demo.queue().launch(allowed_paths=["."])