Upload app.py

app.py

@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+"""Copy of assessment3_Elina_Hemink.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/1xhBZL_ztniX37QTt8SK_mV7nZKO_UrwW
+
+## Create embeddings of the email dataset and store in a chromadb database
+"""
+
+!pip install chromadb
+import chromadb
+from chromadb.utils import embedding_functions
+import pandas as pd
+import email
+from sklearn.model_selection import train_test_split
+
+from google.colab import drive
+drive.mount('/content/drive')
+
+# Loading email.csv dataset
+emails = pd.read_csv('/content/drive/MyDrive/emails.csv')
+print(emails.head())
+
+# What a message looks like
+print(emails['message'][0])
+
+# Getting the content of the emails and saving to a list
+content_text = []
+for item in emails.message:
+  text = email.message_from_string(item)
+  message = (text.get_payload())
+  cleaned_message = message.replace("\n","").replace("\r","").replace("> >>> > >","")
+  content_text.append(cleaned_message)
+
+# Checking content of emails (first 5 items)
+print(content_text[:5])
+
+# Taking a sample of the dataset
+train, test = train_test_split(content_text, train_size = 0.01) # Dataset is too large to complete embedding step
+
+print(train[:5])
+print(len(train))
+
+# Setting up ids for ChromaDB collections
+ids = []
+for i in range(len(train)):
+  id = 'id'+str(i+1)
+  ids.append(id)
+
+# Creating collection
+client = chromadb.Client()
+collection = client.create_collection(name="Enron_emails")
+collection.add(
+    documents = train,
+    ids = ids
+)
+
+"""## Fine-tune a Language Model on the Dataset"""
+
+!pip install transformers[torch] accelerate -U
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments
+
+# Load pre-trained GPT2 tokenizer and model
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+model = GPT2LMHeadModel.from_pretrained('gpt2')
+
+# Tokenize the dataset
+tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+tokenized_emails = tokenizer(train, truncation=True, padding=True)
+
+# Extract token IDs from BatchEncoding object
+token_ids_list = tokenized_emails['input_ids']
+
+# Save token IDs to a text file
+with open('tokenized_emails.txt', 'w') as f:
+    for token_ids in token_ids_list:
+        f.write(' '.join(map(str, token_ids)) + '\n')
+
+# Initialize TextDataset with the file path
+dataset = TextDataset(tokenizer=tokenizer, file_path = 'tokenized_emails.txt', block_size=128)
+
+# Define data collator
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# Define training arguments
+training_args = TrainingArguments(
+    output_dir='./output',
+    num_train_epochs=3,
+    per_device_train_batch_size=8,
+)
+
+# Initialize Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    data_collator=data_collator,
+    train_dataset=dataset,
+)
+
+# Fine-tune the model
+trainer.train()
+
+# Save the fine-tuned model
+model.save_pretrained("/fine_tuned_model")
+tokenizer.save_pretrained("/fine_tuned_model")
+
+"""## Create a Gradio Interface"""
+
+!pip install gradio
+import gradio as gr
+
+model_dir= "/fine_tuned_model"
+tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
+model = GPT2LMHeadModel.from_pretrained(model_dir)
+
+# Load chromadb collection to pass as context
+documents = collection.get(["documents"])
+
+# Define function to answer questions using the fine-tuned model and ChromaDB collection
+def answer_question(question):
+    # Concatenate document contents to create context
+    context = " ".join(doc["content"] for doc in documents)
+
+    # Append question to the context
+    input_text = f"Question: {question} Context: {context} Answer:"
+
+    # Generate answer using the model
+    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    generated = model.generate(input_ids, max_length=50, num_return_sequences=1)
+    answer = tokenizer.decode(generated[0], skip_special_tokens=True)
+
+
+gr.Interface(fn=answer_question, inputs="text", outputs="text").launch()
+
+"""## Deploy the Gradio Interface in a Huggingface Space"""
+