# https://www.kaggle.com/datasets/wcukierski/enron-email-dataset from google.colab import drive drive.mount('/content/drive') # libraries #!pip install transformers --upgrade #!pip install gradio #!pip install datasets #!pip install huggingface-hub #!pip install chromadb #!pip install accelerate==0.21.0 #!pip install transformers[torch] #!pip install git+https://github.com/huggingface/accelerate.git import pandas as pd import numpy as np from transformers import AutoModel from sklearn.model_selection import train_test_split from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline import gradio as gr import chromadb from datasets import Dataset from transformers import Trainer, TrainingArguments from transformers import AutoModelForMaskedLM, DataCollatorForLanguageModeling from transformers import TextDataset, DataCollatorForLanguageModeling #from transformers import TrainingArguments, Trainer #from transformers import pipeline from transformers import AutoTokenizer, AutoModelForSeq2SeqLM file_path = '/content/drive/MyDrive/emails.csv' df = pd.read_csv(file_path) df_columns = df.columns print(df.head(10)) messages_df = df['message'] #extract message column print(messages_df.head()) print(type(messages_df)) # Extract 1% of the content as test set so that instead of 500,000 emails 5,000 are being used as a sample. (Kept changing test size to stop colab crashing.) emails_train, emails_test = train_test_split(messages_df, test_size=0.000008, random_state=42) print(emails_test) print(type(emails_test)) pd.set_option('display.max_colwidth', None) #check content print(emails_test.head()) #first 5 rows print(type(emails_test)) # Embeddings import os # Define maximum sequence length max_seq_length = 512 # Truncate or pad sequences to the maximum length truncated_emails_test = [email[:max_seq_length] for email in emails_test] tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") model = AutoModel.from_pretrained("bert-base-uncased") embeddings_pipeline = pipeline('feature-extraction', model=model, tokenizer=tokenizer) embeddings = embeddings_pipeline(truncated_emails_test) print(type(embeddings)) #print(embeddings[:5]) #cannot see embeddings like this # to see the embeddings # Save each embedding to a separate file for i, emb in enumerate(embeddings): np.save(f"embedding_{i}.npy", emb) # Load each embedding from its corresponding file loaded_embeddings = [] for i in range(len(embeddings)): emb = np.load(f"embedding_{i}.npy") loaded_embeddings.append(emb) for i, emb in enumerate(loaded_embeddings): print(f"Embedding {i}:") print(emb) import chromadb chroma_client = chromadb.Client() collection = chroma_client.create_collection(name="michelletest") # Extract the embeddings from the nested list extracted_embeddings = [embedding[0][0] for embedding in embeddings] # Add embeddings to the ChromaDB collection collection.add( embeddings=extracted_embeddings[:5], # Add the first 5 embeddings documents=emails_test.tolist()[:5], # Add the first 5 documents metadatas=[{"source": "emails_test"} for _ in range(5)], # Metadata for the first 5 documents ids=[f"id{i}" for i in range(5)] # ID for the first 5 documents ) collection.count() #check how many in the database # Retrieve the first 2 entries from the ChromaDB database to check that it worked properly collection.get() # Convert the Series to a DataFrame emails_test_df = emails_test.to_frame() # Print the column names of the DataFrame print(emails_test_df.columns) print(emails_test_df['message']) #checking content of messsages for fine tuning the model print(emails_test_df['message'].head()) # Print the column names of the DataFrame print(emails_test_df.columns) num_entries = emails_test_df.shape[0] print("Number of entries in emails_test_df:", num_entries) # Extract 1% of the content as test set so that instead of 500,000 emails 5,000 are being used as a sample; 60 used in the end emails_train, emails_test2 = train_test_split(messages_df, test_size=0.00001, random_state=42) print(emails_test2) print(type(emails_test2)) num_entries2=emails_test2.shape[0] print("number of",num_entries2) # Convert pandas Series to a list of strings text_list = emails_test_df['message'].tolist() # Verify the type and content print(type(text_list)) print(text_list[:5]) # Print the first 5 entries as an example print(text_list[:5]) print(text_list) print(text_list[2]) #to see the content of an average mail to know what to clean up def remove_sections(email): #clean email of content that is not useful """Remove sections including original message, from, sent, to, subject line, and additional headers.""" sections_to_remove = [ "----- Original Message -----", "From:", "Sent:", "To:", "CC:", "Subject:", "Message-ID:", "Date:", "Mime-Version:", "Content-Type:", "Content-Transfer-Encoding:", "X-cc:", "X-bcc:", "X-Folder:", "X-Origin:", "X-FileName:", "-----Original Message-----" ] for section in sections_to_remove: email = [line for line in email if section not in line] return email # Remove sections from each email in the list cleaned_text_list = [remove_sections(email.split("\n")) for email in text_list] # Print out the cleaned emails to see if content looks ok for cleaned_email in cleaned_text_list: print("\n".join(cleaned_email)) print("=" * 50) # Separate each cleaned email for better readability #fine tune language model # Define the pre-trained model name (bart-base) model_name = "facebook/bart-base" # Load the tokenizer for bart-base tokenizer = AutoTokenizer.from_pretrained(model_name) # Function to preprocess text_list for training def prepare_data(text_list): # Tokenize the text with padding and truncation (BART handles these well) inputs = tokenizer(text_list, padding="max_length", truncation=True) # Copy the input IDs for labels (desired output during training) labels = inputs.input_ids.copy() # Create a Dataset object from the preprocessed data return Dataset.from_dict({"input_ids": inputs["input_ids"], "labels": labels}) """Preprocesses text data for training the BART model. Args: text_list: A list of strings containing the text data. Returns: A Dataset object containing the preprocessed data. """ # Prepare your training data from the text list train_data = prepare_data(text_list) # Define the fine-tuning model (BART for sequence-to-sequence tasks) model = AutoModelForSeq2SeqLM.from_pretrained(model_name) # Training hyperparameters (adjust as needed) batch_size = 8 learning_rate = 2e-5 num_epochs = 3 from transformers import Trainer # Define the Trainer object for training management trainer = Trainer( model=model, args=TrainingArguments( output_dir="./results", # Output directory for checkpoints etc. overwrite_output_dir=True, per_device_train_batch_size=batch_size, learning_rate=learning_rate, num_train_epochs=num_epochs, ), train_dataset=train_data, ) # Start the fine-tuning process trainer.train() # Save the fine-tuned model and tokenizer model.save_pretrained("./fine-tuned_bart") tokenizer.save_pretrained("./fine-tuned_bart") print("Fine-tuning completed! Model saved in ./fine-tuned_bart") # Fine-tuning completed! Model saved in ./fine-tuned_bart # i used a very small amount of input so that colab stopped crashing import gradio as gr from transformers import BartForQuestionAnswering, BartTokenizer # Load the fine-tuned BART model model = BartForQuestionAnswering.from_pretrained("./fine-tuned_bart") tokenizer = BartTokenizer.from_pretrained("./fine-tuned_bart") # Function to answer questions def answer_question(question): inputs = tokenizer.encode_plus(question, return_tensors="pt", max_length=512, truncation=True) input_ids = inputs["input_ids"].tolist()[0] answer_start_scores, answer_end_scores = model(**inputs) answer_start = torch.argmax(answer_start_scores) answer_end = torch.argmax(answer_end_scores) + 1 answer = tokenizer.decode(input_ids[answer_start:answer_end]) return answer # Create Gradio interface iface = gr.Interface( fn=answer_question, inputs="text", outputs="text", title="Question Answering Model", description="Enter a question to get the answer." ) # Launch the interface iface.launch()