final.py

# -*- coding: utf-8 -*-
"""AIE3final.py
______
Automated Grading System for AIE3 Final Project
______
"""

# Import necessary libraries
import logging
import sys
import os
import re
import zipfile
import tempfile
from typing import List, Dict, Tuple
from dotenv import load_dotenv
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from langchain_core.messages import AIMessage
from langchain_openai import OpenAIEmbeddings
from sentence_transformers import SentenceTransformer
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance, PointStruct, ScoredPoint
from docx import Document as DocxDocument
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import getpass
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import openai
import json
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import chainlit as cl
import asyncio
import zipfile

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
openai.api_key = OPENAI_API_KEY

# Set up logging
logging.basicConfig(level=logging.INFO)
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)

def unzip_file(file_path: str, output_dir: str):
    with zipfile.ZipFile(file_path, 'r') as zip_ref:
        for member in zip_ref.namelist():
            if not member.startswith('__MACOSX/'):
                zip_ref.extract(member, output_dir)

def read_pdf(file_path: str) -> List[Document]:
    loader = PyMuPDFLoader(file_path)
    return loader.load()

def read_docx(file_path: str) -> Document:
    doc = DocxDocument(file_path)
    text = "\n".join([p.text for p in doc.paragraphs])
    return Document(page_content=text, metadata={"source": file_path})

def read_files_from_directory(directory: str) -> List[Document]:
    documents = []
    for root, _, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            if os.path.basename(file_path).startswith('~$'):
                continue  # Skip temporary files
            if file_path.endswith('.docx'):
                documents.append(read_docx(file_path))
            elif file_path.endswith('.pdf'):
                documents.extend(read_pdf(file_path))
    return documents

def extract_json(message: AIMessage) -> List[dict]:
    text = message.content
    pattern = r"```json(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    try:
        return [json.loads(match.strip()) for match in matches]
    except Exception:
        raise ValueError(f"Failed to parse: {message}")
    
qa_chat_model = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0
)

ref_prompt = f"""
You are given a reference documents. The document contains a mix of instructions, guides, questions, and answers.
Your task is to go through the reference document and extract questions and answers from the document step-by-step.
Use the keyword 'Question #' to identify the start of each question.
Retain the following words until the 'Answer:' as the question.
Use the keyword 'Answer:' to identify the start of each answer.
Retain the follwing words until the 'Question:' as the answer, until the end of the document.
Remove any white spaces such as carriage returns.
Return the question-answer pairs as a key-value pair as Dict type.
---

Reference Document Content:
{{source}}

Please extract the question-answer pairs and return them as JSON.
"""

ref_prompt_template = ChatPromptTemplate.from_template(ref_prompt)
ref_generation_chain = ref_prompt_template | qa_chat_model

student_prompt = f"""
You are given a student assignment document. The document may contain a mix of instructions, guides, questions, and answers.
Your task is to go through the student document and extract answers to questions from the document step-by-step.
Use the reference document as a guide.
Use the keyword 'Question #' to identify each question.
Then for its associated values, search the student document for the answer.
If you do not see any answer in the student document, return 'No answer found'.
Do not make up any answer.
Remove any white spaces such as carriage returns.
Return the original question and the student answer pairs as a key-value pair as Dict type.
---

Reference Content:
{{source}}

Student Content:
{{student}}

Please extract the question-answer pairs and return them as JSON.
"""

student_prompt_template = ChatPromptTemplate.from_template(student_prompt)
student_response_chain = student_prompt_template | qa_chat_model

def split_documents(documents: List[Document]) -> List[Document]:
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100,
        length_function=len,
        is_separator_regex=False
    )
    split_docs = text_splitter.split_documents(documents)
    total_tokens = sum(len(doc.page_content) for doc in split_docs)  # Approximate token count
    return split_docs, total_tokens

def generate_embeddings(docs: List[Document]) -> List[List[float]]:
    embeddings_model = OpenAIEmbeddings(model="text-embedding-3-small")
    embeddings = embeddings_model.embed_documents([doc.page_content for doc in docs])
    total_tokens = sum(len(doc.page_content) for doc in docs)  # Approximate token count
    return embeddings, total_tokens

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

def process_student(documents, reference):
    test_doc = documents[0]
    student_result = student_response_chain.invoke({"source": reference.keys(),"student": test_doc })
    student_gen_tokens = student_result.usage_metadata["total_tokens"]
    student_result = dict(extract_json(student_result)[0])
    return student_result, student_gen_tokens

def process_reference(reference_document):
    result = ref_generation_chain.invoke({"source": reference_document})
    ref_gen_tokens = result.usage_metadata["total_tokens"]
    reference = dict(extract_json(result)[0])
    
    answers = {}
    for key in reference:
        if key.startswith('Question'):
            question_number = key.split('#')[1]
            answer_key = f'Answer #{question_number}'
            answers[key] = reference[answer_key]
    
    return reference, answers, ref_gen_tokens

def split_docs(answers, student_result):
    split_reference_docs, ref_tokens = {}, 0
    split_student_docs, student_tokens = {}, 0
    for key, value in answers.items():
        split_docs, tokens = split_documents([Document(page_content=value)])
        split_reference_docs[key] = split_docs
        ref_tokens += tokens
    
    for key, value in student_result.items():
        split_docs, tokens = split_documents([Document(page_content=value)])
        split_student_docs[key] = split_docs
        student_tokens += tokens
    
    reference_embeddings = {key: generate_embeddings(value)[0] for key, value in split_reference_docs.items()}
    student_embeddings = {key: generate_embeddings(value)[0] for key, value in split_student_docs.items()}
    
    return reference_embeddings, student_embeddings, ref_tokens, student_tokens

def compute_cosine_similarity(reference_embeddings: dict, student_embeddings: dict) -> float:
    similarity_results = {}
    for key in reference_embeddings.keys():
        if key not in student_embeddings:
            similarity_results[key] = 0
            continue
        reference_vector = np.array(reference_embeddings[key]).reshape(1, -1)
        student_vector = np.array(student_embeddings[key]).reshape(1, -1)
        if reference_vector.shape[1] != student_vector.shape[1]:
            min_dim = min(reference_vector.shape[1], student_vector.shape[1])
            reference_vector = reference_vector[:, :min_dim]
            student_vector = student_vector[:, :min_dim]
        similarity = cosine_similarity(reference_vector, student_vector)[0][0]
        similarity_results[key] = similarity
    
    total_similarity = sum(similarity_results.values())
    num_questions = len(similarity_results)
    average_similarity = total_similarity / num_questions if num_questions else 0
    
    return average_similarity


def llm_similarity(answers, student_result):
    score_prompt = f"""
    You are given two dictionaries representing instructor solution and student answers.
    Your task is to go through each question to grade the correctness of student answer.
    Use the keyword 'Question #' to identify each question.
    Then for its associated values, compare student answer against the instructor answer.
    If the instructor answer has numerical values, check to make sure the student answer has the same number,
    whether it is expressed in numbers or text.
    If you do not see any answer in the student answer, assign score 0 for that answer.
    For student answer that is similar to instructor, assign a full score of 1.
    If the student answer is similar enough, assign a partial score of 0.5.
    Otherwise, assign a score of 0.
    Return the original question and the student score pairs as a key-value pair as Dict type.
    ---

    Reference Content:
    {{source}}

    Student Content:
    {{student}}

    Please extract the question-answer pairs and return them as JSON.
    """

    score_prompt_template = ChatPromptTemplate.from_template(score_prompt)
    student_score_chain = score_prompt_template | qa_chat_model
    
    student_score = student_score_chain.invoke({"source": answers, "student": student_result })
    llm_score_tokens = student_score.usage_metadata["total_tokens"]
    student_score = dict(extract_json(student_score)[0])
    
    total_score = sum(student_score.values())
    num_questions = len(student_score)
    average_score = total_score / num_questions if num_questions else 0
    
    return average_score, llm_score_tokens

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)
    student_total_tokens = student_gen_tokens + student_tokens
    ref_total_tokens = ref_gen_tokens + ref_tokens
    
    average_similarity = compute_cosine_similarity(reference_embeddings, student_embeddings)
    average_score, llm_score_tokens = llm_similarity(answers, student_result)
    llm_total_tokens = ref_gen_tokens + student_gen_tokens + llm_score_tokens
    
    return average_similarity, average_score, ref_total_tokens, student_total_tokens, llm_total_tokens

async def process_grading():
    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():
    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

    # 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, uploaded_file_name

    if message.content.lower() == 'yes' and not user_wants_to_continue:
        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
        await cl.Message(content="Do you want to continue? (yes/no)").send()
    
    elif user_wants_to_continue:
        if message.content.lower() == 'yes':
            user_wants_to_continue = False
            uploaded_file_name = None
            await cl.Message(content="Restarting the app...").send()
            await asyncio.sleep(1)
            python = sys.executable
            os.execl(python, python, *sys.argv)
        
        elif message.content.lower() == 'no':
            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)
            python = sys.executable
            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)
        python = sys.executable
        os.execl(python, python, *sys.argv)

    else:
        await cl.Message(content="Please type 'yes' to start processing or 'no' to exit.").send()