Update app.py

app.py

@@ -1,76 +1,72 @@
-import gradio as gr
+import streamlit as st
 import urllib.request
 import PyPDF2
 import re
 import pandas as pd
-from tqdm import tqdm
 
-def extract_text_from_pdf(pdf_file):
+def convert_pdf_to_txt(pdf_file):
     pdf_reader = PyPDF2.PdfFileReader(pdf_file)
-    text = ""
-    for page in range(pdf_reader.numPages):
-        text += pdf_reader.getPage(page).extractText()
+    text = ''
+    for i in range(pdf_reader.numPages):
+        text += pdf_reader.getPage(i).extractText()
     return text
 
-def extract_text_from_txt(txt_file):
-    with open(txt_file, "r") as file:
-        text = file.read()
+def preprocess_text(text):
+    # Preprocess the text data to remove unwanted characters and convert to lowercase
+    text = re.sub(r'[^\w\s]', '', text)
+    text = text.lower()
     return text
 
-def book_to_dataset(file, file_type):
-    if file_type == "pdf":
-        text = extract_text_from_pdf(file)
-    elif file_type == "txt":
-        text = extract_text_from_txt(file)
-    else:
-        raise ValueError("Invalid file type")
-    words = re.findall(r'\w+', text)
-    words_frequency = {}
-    for word in words:
-        words_frequency[word] = words_frequency.get(word, 0) + 1
-    df = pd.DataFrame(list(words_frequency.items()), columns=["Word", "Frequency"])
-    return df
+def download_book(url):
+    response = urllib.request.urlopen(url)
+    book = response.read()
+    return book
 
-def book_to_dataset_progress(file, file_type):
-    if file_type == "pdf":
-        text = extract_text_from_pdf(file)
-    elif file_type == "txt":
-        text = extract_text_from_txt(file)
-    else:
-        raise ValueError("Invalid file type")
-    words = re.findall(r'\w+', text)
-    words_frequency = {}
-    for word in tqdm(words, desc="Converting..."):
-        words_frequency[word] = words_frequency.get(word, 0) + 1
-    df = pd.DataFrame(list(words_frequency.items()), columns=["Word", "Frequency"])
-    return df
+def upload_book():
+    uploaded_file = st.file_uploader("Choose a book file", type=["pdf", "txt"])
+    if uploaded_file is None:
+        return None
+    return uploaded_file.read()
 
-def book_converter(inputs):
-    if inputs[1] == "URL":
-        url = inputs[0]
-        file_name = url.split("/")[-1]
-        urllib.request.urlretrieve(url, file_name)
-        file = file_name
-        file_type = file_name.split(".")[-1]
-    else:
-        file = inputs[0]
-        file_type = inputs[2].split(".")[-1]
-    return book_to_dataset_progress(file, file_type)
+def main():
+    st.set_page_config(page_title="Book to Dataset Converter", page_icon=":book:", layout="wide")
+    st.title("Book to Dataset Converter")
+    st.write("This app allows you to convert a book to a dataset that can be used to train AI models.")
 
-inputs = (
-    [
-        gr.inputs.Textbox(lines=1, default="Enter URL or choose file", element_type="url"),
-        gr.inputs.Radio(["URL", "File"], default="URL"),
-        gr.inputs.FileUploader(upload_label="Choose file", clear_label="Clear file",)
-    ]
-)
+    source = st.sidebar.radio("Select source of book", ("URL", "Upload"))
 
-interface = gr.Interface(
-    book_converter,
-    inputs,
-    gr.outputs.Dataframe(),
-    title="Book to Dataset Converter",
-    description="Convert a book in pdf or txt format to a dataset that can be used to train AI models."
-)
+    if source == "URL":
+        url = st.sidebar.text_input("Enter URL of book")
+        if st.button("Convert"):
+            with st.spinner("Downloading book..."):
+                book = download_book(url)
+                if book is None:
+                    st.error("Failed to download book")
+            with st.spinner("Converting book to dataset..."):
+                text = convert_pdf_to_txt(book)
+                text = preprocess_text(text)
+                dataset = pd.DataFrame({'text': [text]})
+                st.write(dataset)
+    else:
+        if st.button("Upload"):
+            uploaded_file = upload_book()
+            if uploaded_file is None:
+                st.error("Failed to upload book")
+            else:
+                if uploaded_file.endswith(b".pdf"):
+                    with st.spinner("Converting book to dataset..."):
+                        text = convert_pdf_to_txt(uploaded_file)
+                        text = preprocess_text(text)
+                        dataset = pd.DataFrame({'text': [text]})
+                        st.write(dataset)
+                elif uploaded_file.endswith(b".txt"):
+                    with st.spinner("Converting book to dataset..."):
+                        text = uploaded_file.decode('utf-8')
+                        text = preprocess_text(text)
+                        dataset = pd.DataFrame({'text': [text]})
+                        st.write(dataset)
+                else:
+                    st.error("Invalid file format. Please upload a book in pdf or txt format.")
 
-interface.launch()
+if name == 'main':
+    main()