Update app.py

app.py

@@ -1,89 +1,95 @@
 import streamlit as st
-from transformers import pipeline
-import os
+from config import get_ai71_api_key
+from file_utils import extract_text_from_pdf as _extract_text_from_pdf, extract_text_from_word as _extract_text_from_word
+from ai_utils import generate_response as _generate_response
 
 # Set the page configuration to wide layout
 st.set_page_config(layout="wide")
 
-# Initialize the summarization pipeline
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+# Define the tabs
+tab1, tab2, tab3, tab4, tab5 = st.tabs(["Introduction", "Text Processing", "Summary & Main Points", "Flashcards", "Results"])
 
-# Define functions for summarizing, extracting key points, and generating flashcards
-def summarize_text(text):
-    summary = summarizer(text, max_length=150, min_length=50, do_sample=False)
-    return summary[0]['summary_text']
-
-def extract_key_points(text):
-    summary = summarizer(text, max_length=60, min_length=20, do_sample=False)
-    return [point['summary_text'] for point in summary]
-
-def generate_flashcards(text):
-    # Assume flashcards are similar to key points in this case
-    return extract_key_points(text)
-
-def process_file(file_path):
-    with open(file_path, 'r', encoding='utf-8') as file:
-        text = file.read()
-    summary = summarize_text(text)
-    key_points = extract_key_points(text)
-    flashcards = generate_flashcards(text)
-    return summary, key_points, flashcards
-
-# Define tabs for the application
-tab1, tab2 = st.tabs(["Text Summarization", "File Upload & Process"])
-
-# Text Summarization Tab
 with tab1:
-    st.title("Text Summarization")
-
-    st.subheader("Enter Your Text")
-    user_text = st.text_area("Input Text", height=300)
-
-    if st.button("Generate Summary"):
-        if user_text:
-            summary = summarize_text(user_text)
-            key_points = extract_key_points(user_text)
-            flashcards = generate_flashcards(user_text)
-
-            st.write("### Summary")
-            st.write(summary)
-
-            st.write("### Key Points")
-            for idx, point in enumerate(key_points, start=1):
-                st.write(f"**Point {idx}:** {point}")
+    st.write("""
+        Welcome to the AI71 Text Processing Tool! This application helps you process text files by summarizing content,
+        extracting main points, and creating flashcards. Follow the instructions to upload your files and generate content.
+    """)
 
-            st.write("### Flashcards")
-            for idx, card in enumerate(flashcards, start=1):
-                st.write(f"**Flashcard {idx}:** {card}")
-        else:
-            st.warning("Please enter text to generate a summary.")
-
-# File Upload & Process Tab
 with tab2:
-    st.title("Upload and Process File")
-
-    uploaded_file = st.file_uploader("Choose a text file", type=["txt"])
-
-    if uploaded_file is not None:
-        # Save uploaded file temporarily
-        file_path = os.path.join("uploads", uploaded_file.name)
-        with open(file_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-
-        # Process the file
-        summary, key_points, flashcards = process_file(file_path)
-
-        # Display the summary, key points, and flashcards
-        st.write("### Summary")
-        st.write(summary)
-
-        st.write("### Key Points")
-        for idx, point in enumerate(key_points, start=1):
-            st.write(f"**Point {idx}:** {point}")
-
-        st.write("### Flashcards")
-        for idx, card in enumerate(flashcards, start=1):
-            st.write(f"**Flashcard {idx}:** {card}")
-
-        # Clean up the temporary file
-        os.remove(file_path)
+    AI71_API_KEY = get_ai71_api_key()
+    st.title("Text Processing")
+
+    # Initialize session state variables
+    session_vars = ['uploaded_text', 'summary', 'main_points', 'flashcards', 'results']
+    for var in session_vars:
+        if var not in st.session_state:
+            st.session_state[var] = ""
+
+    # Caching file extraction functions
+    @st.cache_data
+    def extract_text_from_pdf(file):
+        return _extract_text_from_pdf(file)
+    
+    @st.cache_data
+    def extract_text_from_word(file):
+        return _extract_text_from_word(file)
+
+    # Caching API call function
+    @st.cache_data
+    def generate_response(system_message, user_message):
+        return _generate_response(system_message, user_message)
+
+    st.subheader("Upload Text File")
+
+    file = st.file_uploader("Upload a text file (PDF or Word)", type=["pdf", "docx", "doc"])
+    if file is not None:
+        if file.type == "application/pdf":
+            st.session_state.uploaded_text = extract_text_from_pdf(file)
+        elif file.type in ["application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/msword"]:
+            st.session_state.uploaded_text = extract_text_from_word(file)
+        st.text_area("Uploaded Text", value=st.session_state.uploaded_text, height=200, key="uploaded_text")
+
+with tab3:
+    st.title("Summary & Main Points")
+
+    if st.session_state.uploaded_text:
+        summary_prompt = "Summarize the following text into a brief summary:"
+        main_points_prompt = "Extract the main points from the following text:"
+        
+        if st.button("Generate Summary"):
+            st.session_state.summary = generate_response(summary_prompt, st.session_state.uploaded_text)
+            st.text_area("Summary", value=st.session_state.summary, height=200, key="summary")
+
+        if st.button("Extract Main Points"):
+            st.session_state.main_points = generate_response(main_points_prompt, st.session_state.uploaded_text)
+            st.text_area("Main Points", value=st.session_state.main_points, height=200, key="main_points")
+
+with tab4:
+    st.title("Flashcards")
+
+    if st.session_state.uploaded_text:
+        flashcards_prompt = "Create flashcards with questions and answers based on the following text:"
+
+        if st.button("Generate Flashcards"):
+            st.session_state.flashcards = generate_response(flashcards_prompt, st.session_state.uploaded_text)
+            st.text_area("Flashcards", value=st.session_state.flashcards, height=200, key="flashcards")
+
+with tab5:
+    st.title("Results")
+    st.subheader("Summary")
+    st.text_area("Summary", value=st.session_state.summary, height=200, key="summary_results")
+    
+    st.subheader("Main Points")
+    st.text_area("Main Points", value=st.session_state.main_points, height=200, key="main_points_results")
+    
+    st.subheader("Flashcards")
+    st.text_area("Flashcards", value=st.session_state.flashcards, height=200, key="flashcards_results")
+
+    st.subheader("Download Results")
+    results = f"Summary:\n{st.session_state.summary}\n\nMain Points:\n{st.session_state.main_points}\n\nFlashcards:\n{st.session_state.flashcards}"
+    st.download_button(
+        label="Download Results",
+        data=results,
+        file_name="results.txt",
+        mime="text/plain"
+    )