zip

final.py

@@ -34,6 +34,7 @@ import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 import chainlit as cl
 import asyncio
+import zipfile
 
 # Load environment variables
 load_dotenv()
@@ -47,6 +48,7 @@ logger = logging.getLogger(__name__)
 # Define constants
 REFERENCE_DOCUMENT_PATH = './Excel Review.pdf'
 UPLOAD_FOLDER = './uploads'
+TEMP_DIR = "./temp"
 
 # Ensure the upload folder exists
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
@@ -154,9 +156,9 @@ def generate_embeddings(docs: List[Document]) -> List[List[float]]:
     total_tokens = sum(len(doc.page_content) for doc in docs)  # Approximate token count
     return embeddings, total_tokens
 
-def prepare_files():
-    unzip_file('./uploads/fall23_small.zip', './temp')
-    documents = read_files_from_directory('./temp/fall23_small')
+def prepare_files(zip_file_name: str):
+    unzip_file(os.path.join(UPLOAD_FOLDER, zip_file_name), TEMP_DIR)
+    documents = read_files_from_directory(os.path.join(TEMP_DIR, os.path.splitext(zip_file_name)[0]))
     reference_document = read_pdf(REFERENCE_DOCUMENT_PATH)
     return documents, reference_document
 
@@ -258,8 +260,8 @@ def llm_similarity(answers, student_result):
     
     return average_score, llm_score_tokens
 
-def process_data() -> Tuple[float, float, int, int, int]:
-    documents, reference_document = prepare_files()
+def process_data(zip_file_name: str) -> Tuple[float, float, int, int, int]:
+    documents, reference_document = prepare_files(zip_file_name)
     reference, answers, ref_gen_tokens = process_reference(reference_document)
     student_result, student_gen_tokens = process_student(documents, reference)
     reference_embeddings, student_embeddings, ref_tokens, student_tokens = split_docs(answers, student_result)
@@ -273,33 +275,61 @@ def process_data() -> Tuple[float, float, int, int, int]:
     return average_similarity, average_score, ref_total_tokens, student_total_tokens, llm_total_tokens
 
 async def process_grading():
-    average_similarity, average_score, ref_total_tokens, student_total_tokens, llm_total_tokens = process_data()
-    
-    await cl.Message(content=f"Total tokens used for reference documents: {ref_total_tokens}").send()
-    await cl.Message(content=f"Total tokens used for student documents: {student_total_tokens}").send()
-    await cl.Message(content=f"Total tokens used by LLM: {llm_total_tokens}").send()
-    await cl.Message(content=f"Score: {average_similarity}").send()
-    await cl.Message(content=f"Average Score: {average_score}").send()
+    global uploaded_file_name
+    if uploaded_file_name:
+        try:
+            # Process the uploaded ZIP file
+            average_similarity, average_score, ref_total_tokens, student_total_tokens, llm_total_tokens = process_data(uploaded_file_name)
+
+            # Send results
+            await cl.Message(content=f"Processing complete. Results:\n"
+                                     f"Average Similarity: {average_similarity:.2f}\n"
+                                     f"Average Score: {average_score:.2f}\n"
+                                     f"Reference Total Tokens: {ref_total_tokens}\n"
+                                     f"Student Total Tokens: {student_total_tokens}\n"
+                                     f"LLM Total Tokens: {llm_total_tokens}").send()
+        except Exception as e:
+            await cl.Message(f"An error occurred while processing the zip file: {str(e)}").send()
+    else:
+        await cl.Message("No file has been uploaded yet. Please upload a ZIP file first.").send()
+
+user_wants_to_continue = False
 
 @cl.on_chat_start
-async def start_chat():
-    await cl.Message(content="Do you want to proceed with the grading? (yes/no)").send()
+async def start():
+    global uploaded_file_name
+    files = None
 
+    # Wait for the user to upload a file
+    while files is None:
+        files = await cl.AskFileMessage(
+            content="Please upload a zip file to begin!", accept={"application/zip": [".zip"]}
+        ).send()
 
+    zip_file = files[0]  # Assuming only one file is uploaded
+    file_path = os.path.join(UPLOAD_FOLDER, zip_file.name)
+    uploaded_file_name = zip_file.name
 
-# Define a global flag to track the processing state
-user_wants_to_continue = False
+    # Move the uploaded file to the desired location
+    os.rename(zip_file.path, file_path)
+
+    # Let the user know that the system is ready
+    await cl.Message(content=f"`{zip_file.name}` uploaded successfully!").send()
+
+    # Ask if the user wants to proceed with grading
+    await cl.Message(content="Do you want to proceed with the grading? (yes/no)").send()
 
 @cl.on_message
 async def on_message(message: cl.Message):
-    global user_wants_to_continue
+    global user_wants_to_continue, uploaded_file_name
 
     if message.content.lower() == 'yes' and not user_wants_to_continue:
-        # Start processing
-        processing_message = cl.Message(content="Processing files...")
-        await processing_message.send()  # Send the message immediately
-        await asyncio.sleep(0.5)  # Short delay to ensure the message is displayed
-        await process_grading()
+        if uploaded_file_name:
+            # Start processing
+            processing_message = cl.Message(content="Processing files...")
+            await processing_message.send()
+            await asyncio.sleep(0.5)
+            await process_grading()
         
         # Ask user if they want to continue after processing is done
         user_wants_to_continue = True
@@ -307,27 +337,29 @@ async def on_message(message: cl.Message):
     
     elif user_wants_to_continue:
         if message.content.lower() == 'yes':
-            user_wants_to_continue = False  # Reset the flag
+            user_wants_to_continue = False
+            uploaded_file_name = None
             await cl.Message(content="Restarting the app...").send()
-            await asyncio.sleep(1)  # Give time for the message to be sent
+            await asyncio.sleep(1)
             python = sys.executable
-            os.execl(python, python, *sys.argv)  # Restart the app
+            os.execl(python, python, *sys.argv)
         
         elif message.content.lower() == 'no':
-            user_wants_to_continue = False  # Reset the flag
+            user_wants_to_continue = False
+            uploaded_file_name = None
             await cl.Message(content="Okay, thank you for using the grading app. Restarting...").send()
-            await asyncio.sleep(1)  # Give time for the message to be sent
+            await asyncio.sleep(1)
             python = sys.executable
-            os.execl(python, python, *sys.argv)  # Restart the app
+            os.execl(python, python, *sys.argv)
         
         else:
             await cl.Message(content="Invalid response. Please type 'yes' or 'no'.").send()
     
     elif message.content.lower() == 'no':
         await cl.Message(content="Okay, thank you for using the grading app. Restarting...").send()
-        await asyncio.sleep(1)  # Give time for the message to be sent
+        await asyncio.sleep(1)
         python = sys.executable
-        os.execl(python, python, *sys.argv)  # Restart the app
+        os.execl(python, python, *sys.argv)
 
     else:
         await cl.Message(content="Please type 'yes' to start processing or 'no' to exit.").send()