Update app.py

app.py

@@ -1,89 +1,85 @@
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
-from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
-
-# =============================
-# Словарь доступных моделей
-# =============================
-model_options = {
-    "Русский → Английский": "Helsinki-NLP/opus-mt-ru-en",
-    "Английский → Русский": "Helsinki-NLP/opus-mt-en-ru",
-    "Казахский → Русский": "Helsinki-NLP/opus-mt-kk-ru",
-    "Русский → Казахский": "Helsinki-NLP/opus-mt-ru-kk",
-}
-
-# =============================
-# Кэш моделей для ускорения
-# =============================
-model_cache = {}
-
-def load_translation_model(model_name):
-    if model_name in model_cache:
-        return model_cache[model_name]
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    model_cache[model_name] = (model, tokenizer)
-    return model, tokenizer
-
-# =============================
-# Функция перевода
-# =============================
-def translate_interface(direction, text):
-    if not text or len(text.strip()) == 0:
-        return "Введите текст для перевода!"
-    model_name = model_options[direction]
-    model, tokenizer = load_translation_model(model_name)
-    batch = tokenizer([text], return_tensors="pt", padding=True)
-    gen = model.generate(**batch)
-    result = tokenizer.batch_decode(gen, skip_special_tokens=True)[0]
-    return result
-
-# =============================
-# Функция BLEU
-# =============================
-def bleu_interface(direction, original, reference):
-    if not original.strip() or not reference.strip():
-        return "Введите текст и эталон!"
-    translation = translate_interface(direction, original)
-    smoothie = SmoothingFunction().method4
-    ref_tokens = [reference.split()]
-    hyp_tokens = translation.split()
-    bleu = sentence_bleu(ref_tokens, hyp_tokens, smoothing_function=smoothie)
-    return f"BLEU: {bleu:.4f}"
-
-# =============================
-# Gradio UI
-# =============================
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", neutral_hue="slate", secondary_hue="pink")) as demo:
-    gr.Markdown("""
-    # 🌙 Мощный Машинный Перевод
-    Тёмная тема, стильный UI, BLEU, множество моделей.
-    """)
-
-    with gr.Row():
-        direction = gr.Dropdown(list(model_options.keys()), label="Направление перевода")
-
-    with gr.Row():
-        input_text = gr.Textbox(lines=4, label="Введите текст")
-        output_text = gr.Textbox(lines=4, label="Перевод")
-
-    translate_btn = gr.Button("🚀 Перевести")
-    translate_btn.click(fn=translate_interface, inputs=[direction, input_text], outputs=output_text)
-
-    # Кнопка копирования перевода
-    copy_btn = gr.Button("📋 Копировать перевод")
-    copy_btn.click(lambda text: text, inputs=output_text, outputs=output_text)
-
-    gr.Markdown("---")
-    gr.Markdown("### 📊 BLEU Оценка качества")
-
-    with gr.Row():
-        original = gr.Textbox(lines=3, label="Исходный текст")
-        reference = gr.Textbox(lines=3, label="Эталонный перевод")
-        bleu_out = gr.Textbox(label="BLEU Score")
-
-    bleu_btn = gr.Button("📈 Посчитать BLEU")
-    bleu_btn.click(fn=bleu_interface, inputs=[direction, original, reference], outputs=bleu_out)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+import re
+
+
+# --------------------------
+# Очистка текста
+# --------------------------
+def clean_text(s):
+    s = s.strip()
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r"\s+([,.!?;:])", r"\1", s)
+    return s
+
+
+# --------------------------
+# Загрузка модели
+# --------------------------
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return tokenizer, model
+
+
+# --------------------------
+# Перевод
+# --------------------------
+def translate_text(text, model_name):
+    text = clean_text(text)
+
+    tokenizer, model = load_model(model_name)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+
+    enc = tokenizer([text], return_tensors="pt", padding=True, truncation=True)
+    enc = {k: v.to(device) for k, v in enc.items()}
+
+    out = model.generate(**enc, max_length=150, num_beams=4)
+    translated = tokenizer.decode(out[0], skip_special_tokens=True)
+
+    return translated
+
+
+# --------------------------
+# Gradio Interface
+# --------------------------
+def build_ui():
+    title = "🌐 Automatic Text Translator"
+    description = """
+    Лёгкий переводчик на HuggingFace Transformers.  
+    Выберите модель → введите текст → получите перевод.
+    """
+
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown(f"# {title}")
+        gr.Markdown(description)
+
+        with gr.Row():
+            inp = gr.Textbox(label="Введите текст", placeholder="Hello, how are you?")
+            out = gr.Textbox(label="Перевод")
+
+        model_name = gr.Dropdown(
+            label="Модель перевода",
+            value="Helsinki-NLP/opus-mt-en-ru",
+            choices=[
+                "Helsinki-NLP/opus-mt-en-ru",
+                "Helsinki-NLP/opus-mt-ru-en",
+                "Helsinki-NLP/opus-mt-en-de",
+                "Helsinki-NLP/opus-mt-en-fr",
+                "Helsinki-NLP/opus-mt-en-kaz"
+            ],
+        )
+
+        btn = gr.Button("Перевести")
+
+        btn.click(translate_text, inputs=[inp, model_name], outputs=out)
+
+    return demo
+
+
+demo = build_ui()
 
 if __name__ == "__main__":
     demo.launch()