import gradio as gr
from transformers import pipeline
from PyPDF2 import PdfReader
import numpy as np
from bark import generate_audio, preload_models
from scipy.io.wavfile import write as write_wav
import torch
import tempfile
import os

# Preload models if needed
preload_models()

def summarize_abstract_from_pdf(pdf_file):
    # Function to extract and summarize the abstract from a PDF
    abstract_string = 'abstract'
    found_abstract = False
    intro_string = 'introduction'
    extracted_text_string = ""

    # Read the PDF and extract text from the first page
    reader = PdfReader(pdf_file)
    text = reader.pages[0].extract_text()

    for line in text.splitlines():
        lower_line = line.lower()
        if lower_line.strip() == abstract_string:
            found_abstract = True
        elif "1" in lower_line.strip() and intro_string in lower_line.strip():
            found_abstract = False

        if found_abstract:
            extracted_text_string += line + " "

    extracted_text_string = extracted_text_string.replace("Abstract", "")

    # Use Hugging Face summarization pipeline
    summarizer = pipeline("summarization", "pszemraj/led-base-book-summary", device=0 if torch.cuda.is_available() else -1)
    summarized_abstract = summarizer(extracted_text_string, min_length=16, max_length=150, no_repeat_ngram_size=3, encoder_no_repeat_ngram_size=3, repetition_penalty=3.5, num_beams=4, early_stopping=True)
    return summarized_abstract[0]['summary_text']

def generate_audio_func(pdf_file):
    text_prompt = summarize_abstract_from_pdf(pdf_file)
    audio_array = generate_audio(text_prompt)

    # Create a temporary WAV file to save the audio
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav_file:
        write_wav(temp_wav_file.name, 22050, (audio_array * 32767).astype(np.int16))
        return temp_wav_file.name

# Define the Gradio interface
demo = gr.Interface(
    fn=generate_audio_func,
    inputs=gr.inputs.File(file_types=["pdf"]),
    outputs=gr.outputs.Audio(type="file"),
    title="PDF to Audio Converter",
    description="Convert text from a PDF file to audio. Upload a PDF file with an abstract to get started."
)

if __name__ == "__main__":
    demo.launch()