Update app.py

app.py

@@ -6,22 +6,6 @@ import pyspark.sql.functions as F
 from llama_cpp import Llama
 from loguru import logger  # Import the logger from loguru
 
-# Create the models directory
-!mkdir -p ./models
-
-# Download the Llama model files
-!wget -O ./models/llama-2-7b-chat.ggmlv3.q8_0.bin https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML/resolve/main/llama-2-7b-chat.ggmlv3.q8_0.bin
-!wget -O ./models/llama-2-7b-chat.ggmlv3.q2_K.bin https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML/resolve/main/llama-2-7b-chat.ggmlv3.q2_K.bin
-
-# download "War and Peace" from Project Gutenberg
-!mkdir -p ./data
-!curl "https://gutenberg.org/cache/epub/2600/pg2600.txt" -o ./data/war_and_peace.txt
-
-
-# Define the Llama models
-MODEL_Q8_0 = Llama(model_path="./models/llama-2-7b-chat.ggmlv3.q8_0.bin", n_ctx=8192, n_batch=512)
-MODEL_Q2_K = Llama(model_path="./models/llama-2-7b-chat.ggmlv3.q2_K.bin", n_ctx=8192, n_batch=512)
-
 # Function to read the text file and create Spark DataFrame
 def create_spark_dataframe(text):
     # Get list of chapter strings
@@ -34,6 +18,15 @@ def create_spark_dataframe(text):
 
     return df
 
+# Read the "War and Peace" text file and create Spark DataFrame
+with open('/content/data/war_and_peace.txt', 'r') as file:
+    text = file.read()
+df_chapters = create_spark_dataframe(text)
+
+# Define the Llama models
+MODEL_Q8_0 = Llama(model_path="./models/llama-2-7b-chat.ggmlv3.q8_0.bin", n_ctx=8192, n_batch=512)
+MODEL_Q2_K = Llama(model_path="./models/llama-2-7b-chat.ggmlv3.q2_K.bin", n_ctx=8192, n_batch=512)
+
 # Function to summarize a chapter using the selected model
 def llama2_summarize(chapter_text, model_version):
     # Choose the model based on the model_version parameter
@@ -74,34 +67,19 @@ def llama2_summarize(chapter_text, model_version):
 
     return summary
 
-# Read the "War and Peace" text file and create Spark DataFrame
-with open('/content/data/war_and_peace.txt', 'r') as file:
-    text = file.read()
-df_chapters = create_spark_dataframe(text)
-
-# Create summaries via Spark
-summaries = (df_chapters
-                .limit(1)
-                .groupby('chapter')
-                .applyInPandas(llama2_summarize, schema='summary string, chapter int')
-                .show(vertical=True, truncate=False)
-            )
-
-# Prompt for the file
-file_path = gr.inputs.File(label="Upload 'War and Peace' text file")
-
-# Choose the model version
-model_version = gr.inputs.Radio(["q8_0", "q2_K"], label="Choose Model Version")
-
 # Define the Gradio interface
 iface = gr.Interface(
     fn=llama2_summarize,
-    inputs=[file_path, model_version],
+    inputs=[
+        gr.inputs.File(label="Upload Text File"),
+        "text",
+    ],  # chapter_text, model_version
     outputs="text",  # Summary text
     live=False,
     capture_session=True,
     title="Llama2 Chapter Summarizer",
-    description="Upload a text file of the novel 'War and Peace', and choose the model version ('q8_0' or 'q2_K') to get a summarized sentence for each chapter.",
+    description="Upload the text file or enter the chapter text and model version ('q8_0' or 'q2_K') to get a summarized sentence.",
 )
 
-iface.launch();
+if __name__ == "__main__":
+    iface.launch();