import gradio as gr
from toshiba import extract_toshiba_data
from bhel import extract_bhel_data

def extract_data(pdf_file, company):
    if company == 'Toshiba':
        return extract_toshiba_data(pdf_file)
    elif company == 'BHEL':
        return extract_bhel_data(pdf_file)
    else:
        raise ValueError("Unsupported company format")

company_options = ['Toshiba', 'BHEL']
interface = gr.Interface(
    fn=extract_data,
    inputs=[gr.File(label="Upload PDF"), gr.Dropdown(choices=company_options, label="Select Company")],
    outputs=gr.File(label="Download Extracted Data as Excel"),
    title="PDF Data Extractor for Toshiba and BHEL",
    description="Upload a PDF file and select the company to extract and format data into an Excel file according to specific requirements."
)

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