Update app.py

app.py

@@ -1,5 +1,5 @@
 from PyPDF2 import PdfReader
-from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
+from transformers import pipeline, MarianMTModel, MarianTokenizer
 import torch
 import gradio as gr
 from gtts import gTTS
@@ -8,12 +8,12 @@ from gtts import gTTS
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Инициализация модели для суммаризации
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0 if device.type == "cuda" else -1)
+summarizer = pipeline("summarization", model="philschmid/bart-large-cnn-samsum", device=0 if device.type == "cuda" else -1)
 
 # Инициализация модели для перевода
-model_translation = T5ForConditionalGeneration.from_pretrained('utrobinmv/t5_translate_en_ru_zh_small_1024')
-model_translation.to(device)
-tokenizer_translation = T5Tokenizer.from_pretrained('utrobinmv/t5_translate_en_ru_zh_small_1024')
+translation_model_name = "Helsinki-NLP/opus-mt-en-ru"
+translation_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
+translation_model = MarianMTModel.from_pretrained(translation_model_name).to(device)
 
 def parse_pdf(pdf_file):
     """Функция для извлечения текста из PDF файла."""
@@ -23,23 +23,24 @@ def parse_pdf(pdf_file):
         extracted_text += page.extract_text() or ""
     return extracted_text
 
-def summarize(text, max_length=1000, min_length=150):
+def summarize(text, max_length=130, min_length=30):
     """Функция для суммаризации текста."""
-    max_length = 1000  # Можно настроить
-    truncated_text = text[:max_length]
-    result = summarizer(truncated_text, max_length=max_length, min_length=min_length, do_sample=False)
-    return result[0]['summary_text']
+    # Разделяем текст на части, если он слишком большой
+    max_chunk_size = 1024
+    chunks = [text[i:i + max_chunk_size] for i in range(0, len(text), max_chunk_size)]
+    summaries = []
+    for chunk in chunks:
+        result = summarizer(chunk, max_length=max_length, min_length=min_length, do_sample=False)
+        summaries.append(result[0]['summary_text'])
+    return " ".join(summaries)
 
 def translate(text):
     """Функция для перевода текста на русский."""
-    prefix = 'translate to ru: '
-    src_text = prefix + text
-
-    input_ids = tokenizer_translation(src_text, return_tensors="pt")
-    generated_tokens = model_translation.generate(**input_ids.to(device))
-
-    result = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True)
-    return result[0]
+    # Токенизация и перевод
+    inputs = translation_tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    translated_tokens = translation_model.generate(**inputs)
+    translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    return translated_text
 
 def text_to_speech(text, language='ru'):
     """Функция для преобразования текста в аудиофайл."""
@@ -48,7 +49,7 @@ def text_to_speech(text, language='ru'):
     tts.save(audio_file)
     return audio_file
 
-def process_pdf(pdf_file):
+def process_pdf(pdf_file, language):
     """Основная функция обработки PDF файла."""
     if not pdf_file:
         return "No input provided."
@@ -61,25 +62,32 @@ def process_pdf(pdf_file):
     summary = summarize(extracted_text)
     print(f"Summary: {summary}")
 
-    # Перевод текста на русский
-    translated_text = translate(summary)
-    print(f"Translated Text: {translated_text}")
+    # Если выбран русский язык, переводим текст
+    if language == "rus":
+        translated_text = translate(summary)
+        print(f"Translated Text: {translated_text}")
+        final_text = translated_text
+        audio_language = 'ru'
+    else:
+        final_text = summary
+        audio_language = 'en'
 
     # Преобразование текста в аудио
-    audio_file = text_to_speech(translated_text)
-    return translated_text, audio_file
+    audio_file = text_to_speech(final_text, language=audio_language)
+    return final_text, audio_file
 
 # Создание Gradio интерфейса
 with gr.Blocks() as demo:
     gr.Markdown("# PDF Summarizer, Translator, and Text-to-Speech")
-    gr.Markdown("Upload a PDF file to summarize, translate to Russian, and convert to audio.")
+    gr.Markdown("Upload a PDF file to summarize, translate (if needed), and convert to audio.")
 
+    language = gr.Radio(choices=["eng", "rus"], label="Output Language", value="rus")
     pdf_input = gr.File(label="Upload PDF File", type="filepath")
-    text_output = gr.Textbox(label="Translated Text", lines=10)
+    text_output = gr.Textbox(label="Processed Text", lines=10)
     audio_output = gr.Audio(label="Generated Audio", type="filepath")
     process_button = gr.Button("Process PDF")
 
-    process_button.click(process_pdf, inputs=pdf_input, outputs=[text_output, audio_output])
+    process_button.click(process_pdf, inputs=[pdf_input, language], outputs=[text_output, audio_output])
 
 # Запуск приложения
 demo.launch(debug=True)