Update app.py

app.py

@@ -0,0 +1,151 @@
+import gradio as gr
+import os
+import json
+from datasets import Dataset
+from transformers import (
+    MarianMTModel, MarianTokenizer,
+    T5ForConditionalGeneration, T5Tokenizer,
+    DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer
+)
+import torch
+
+os.makedirs("models", exist_ok=True)
+
+# ----------- LOAD MODELS -----------
+
+BASE_MODELS = {
+    "MarianMT ru→en": "Helsinki-NLP/opus-mt-ru-en",
+    "MarianMT en→ru": "Helsinki-NLP/opus-mt-en-ru",
+    "T5-small ru→en": "t5-small",
+    "T5-small en→ru": "t5-small"
+}
+
+def load_model(model_id):
+    if "Marian" in model_id:
+        tokenizer = MarianTokenizer.from_pretrained(model_id)
+        model = MarianMTModel.from_pretrained(model_id)
+    else:
+        tokenizer = T5Tokenizer.from_pretrained(model_id)
+        model = T5ForConditionalGeneration.from_pretrained(model_id)
+    return model, tokenizer
+
+# ----------- TRAINING FUNCTION -----------
+
+def train_model(base_model_name, train_file, num_epochs, batch_size):
+
+    # load dataset
+    data = train_file.decode("utf-8").split("\n")
+    pairs = [l.split("\t") for l in data if "\t" in l]
+
+    ds = Dataset.from_dict({
+        "src": [p[0] for p in pairs],
+        "trg": [p[1] for p in pairs]
+    })
+
+    # load pretrained
+    model_id = BASE_MODELS[base_model_name]
+    model, tokenizer = load_model(model_id)
+
+    # preprocess function
+    def preprocess(batch):
+        if "Marian" in base_model_name:
+            inputs = tokenizer(batch["src"], truncation=True, padding="max_length", max_length=128)
+            with tokenizer.as_target_tokenizer():
+                labels = tokenizer(batch["trg"], truncation=True, padding="max_length", max_length=128)
+            inputs["labels"] = labels["input_ids"]
+            return inputs
+        else:  # T5
+            prefix = "translate Russian to English: " if "ru→en" in base_model_name else "translate English to Russian: "
+            inputs = tokenizer(prefix + batch["src"], truncation=True, padding="max_length", max_length=128)
+            with tokenizer.as_target_tokenizer():
+                labels = tokenizer(batch["trg"], truncation=True, padding="max_length", max_length=128)
+            inputs["labels"] = labels["input_ids"]
+            return inputs
+
+    tokenized = ds.map(preprocess, batched=True)
+
+    # training args
+    args = Seq2SeqTrainingArguments(
+        output_dir="models",
+        metric_for_best_model="loss",
+        save_strategy="no",
+        num_train_epochs=num_epochs,
+        per_device_train_batch_size=batch_size,
+        learning_rate=2e-4,
+        logging_steps=5,
+        report_to="none",
+    )
+
+    collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=args,
+        train_dataset=tokenized,
+        data_collator=collator,
+    )
+
+    trainer.train()
+
+    # SAVE
+    save_path = f"models/{base_model_name.replace(' ', '_')}"
+    model.save_pretrained(save_path)
+    tokenizer.save_pretrained(save_path)
+
+    return f"Модель сохранена в {save_path}"
+
+# ----------- TRANSLATION -----------
+
+def translate(text, model_name):
+    model_path = f"models/{model_name.replace(' ', '_')}"
+    if not os.path.exists(model_path):
+        return "Сначала обучите модель."
+
+    if "Marian" in model_name:
+        tokenizer = MarianTokenizer.from_pretrained(model_path)
+        model = MarianMTModel.from_pretrained(model_path)
+    else:
+        tokenizer = T5Tokenizer.from_pretrained(model_path)
+        model = T5ForConditionalGeneration.from_pretrained(model_path)
+
+    if "T5-small" in model_name:
+        prefix = "translate Russian to English: " if "ru→en" in model_name else "translate English to Russian: "
+        input_ids = tokenizer(prefix + text, return_tensors="pt").input_ids
+        out = model.generate(input_ids, max_length=200)
+        return tokenizer.decode(out[0], skip_special_tokens=True)
+
+    else:  # Marian
+        enc = tokenizer([text], return_tensors="pt")
+        out = model.generate(**enc)
+        return tokenizer.decode(out[0], skip_special_tokens=True)
+
+
+# ----------- GRADIO UI -----------
+
+with gr.Blocks() as demo:
+
+    gr.Markdown("# 🚀 Обучение переводчика (MarianMT / T5-small)")
+
+    with gr.Tab("Обучение"):
+        base_model = gr.Dropdown(list(BASE_MODELS.keys()), label="Выберите модель")
+        train_data = gr.File(label="Загрузите тренировочный датасет (формат: src<TAB>tgt)")
+        epochs = gr.Slider(1, 5, value=1, step=1, label="Эпохи")
+        batch = gr.Slider(1, 16, value=4, step=1, label="Батч")
+
+        train_button = gr.Button("Начать обучение")
+        train_output = gr.Textbox(label="Логи")
+
+        train_button.click(
+            train_model,
+            inputs=[base_model, train_data, epochs, batch],
+            outputs=train_output
+        )
+
+    with gr.Tab("Перевод"):
+        model_choice = gr.Dropdown(list(BASE_MODELS.keys()), label="Выберите обученную модель")
+        text = gr.Textbox(lines=5, label="Введите текст")
+        translate_button = gr.Button("Перевести")
+        translation_result = gr.Textbox(label="Перевод")
+
+        translate_button.click(translate, [model_choice, text], translation_result)
+
+demo.launch()