import streamlit as st
from PyPDF2 import PdfReader
from docx import Document
import pandas as pd
from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
from datasets import Dataset

# Function to extract text from PDF
def extract_text_from_pdf(pdf_file):
    reader = PdfReader(pdf_file)
    text = ''
    for page in reader.pages:
        text += page.extract_text()
    return text

# Function to extract text from DOCX
def extract_text_from_docx(docx_file):
    doc = Document(docx_file)
    text = ''
    for para in doc.paragraphs:
        text += para.text
    return text

# Function to extract text from Excel
def extract_text_from_excel(excel_file):
    df = pd.read_excel(excel_file)
    return df.to_string()

# Streamlit UI elements
st.title("Document Text Extractor and LLM Fine-Tuning")

# Upload multiple files of specific types
uploaded_files = st.file_uploader("Upload your files", type=['pdf', 'docx', 'xlsx'], accept_multiple_files=True)

# If the 'Extract Text' button is clicked
if st.button("Extract Text"):
    complete_text = ""
    if uploaded_files is not None:
        # Loop through all the uploaded files and extract text
        for uploaded_file in uploaded_files:
            if uploaded_file.name.endswith('.pdf'):
                complete_text += extract_text_from_pdf(uploaded_file)
            elif uploaded_file.name.endswith('.docx'):
                complete_text += extract_text_from_docx(uploaded_file)
            elif uploaded_file.name.endswith('.xlsx'):
                complete_text += extract_text_from_excel(uploaded_file)

        # Display the extracted text
        st.text_area("Extracted Text", complete_text, height=300)

        # Save the extracted text to a local file
        with open("extracted_text.txt", "w", encoding="utf-8") as file:
            file.write(complete_text)
        st.success("Text saved locally as `extracted_text.txt`!")

        # Option to start fine-tuning on the extracted text
        if st.button("Start Fine-Tuning"):
            # Load the extracted text as a dataset
            data_dict = {"text": [complete_text]}
            dataset = Dataset.from_dict(data_dict)

            # Define the pre-trained GPT-2 model and tokenizer
            model_name = "gpt2"
            tokenizer = GPT2Tokenizer.from_pretrained(model_name)
            model = GPT2LMHeadModel.from_pretrained(model_name)

            # Tokenize the dataset
            def tokenize_function(examples):
                return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

            # Apply tokenization
            tokenized_dataset = dataset.map(tokenize_function, batched=True)

            # Define training arguments
            training_args = TrainingArguments(
                output_dir="./results",
                num_train_epochs=1,  # You can increase this for better results
                per_device_train_batch_size=2,
                save_steps=10_000,
                save_total_limit=2,
                logging_dir='./logs',
                logging_steps=500,
            )

            # Create Trainer
            trainer = Trainer(
                model=model,
                args=training_args,
                train_dataset=tokenized_dataset,
            )

            # Start fine-tuning
            trainer.train()
            trainer.save_model("./fine_tuned_model")

            st.success("Fine-tuning completed and model saved!")

# Option to generate text using the fine-tuned model
st.header("Generate Text Using Fine-Tuned Model")

# Load the fine-tuned model and tokenizer
if st.button("Load Fine-Tuned Model"):
    fine_tuned_model = GPT2LMHeadModel.from_pretrained("./fine_tuned_model")
    fine_tuned_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    st.success("Fine-tuned model loaded successfully!")

# Prompt input for text generation
user_prompt = st.text_input("Enter your prompt for text generation:")

# Generate text using the fine-tuned model
if st.button("Generate Text"):
    if 'fine_tuned_model' in locals() and user_prompt:
        inputs = fine_tuned_tokenizer.encode(user_prompt, return_tensors="pt")
        outputs = fine_tuned_model.generate(inputs, max_length=50, num_return_sequences=1)
        generated_text = fine_tuned_tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.write("Generated Text:", generated_text)
    else:
        st.write("Please load the fine-tuned model and provide a prompt.")