Enron / app.py
user2434's picture
Create app.py
746b99c verified
raw
history blame
No virus
2.56 kB
import pandas as pd
import chromadb
from sklearn.model_selection import train_test_split
from transformers import GPT2Tokenizer, GPT2LMHeadModel, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments, pipeline
import gradio as gr
import email
# loading and preprocessing dataset
emails = pd.read_csv('emails.csv')
def preprocess_email_content(raw_email):
message = email.message_from_string(raw_email).get_payload()
return message.replace("\n", "").replace("\r", "").replace("> >>> > >", "").strip()
content_text = [preprocess_email_content(item) for item in emails['message']]
train_content, _ = train_test_split(content_text, train_size=0.00005) # was unable to load more emails due to capacity constraints
# ChromaDB setup
client = chromadb.Client()
collection = client.create_collection(name="Enron_emails")
collection.add(documents=train_content, ids=[f'id{i+1}' for i in range(len(train_content))])
# model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
# tokenizing and training
tokenized_emails = tokenizer(train_content, truncation=True, padding=True)
with open('tokenized_emails.txt', 'w') as file:
for ids in tokenized_emails['input_ids']:
file.write(' '.join(map(str, ids)) + '\n')
dataset = TextDataset(tokenizer=tokenizer, file_path='tokenized_emails.txt', block_size=128)
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
training_args = TrainingArguments(
output_dir='./output',
num_train_epochs=3,
per_device_train_batch_size=8
)
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=dataset
)
trainer.train()
# saving the model
model.save_pretrained("./fine_tuned_model")
tokenizer.save_pretrained("./fine_tuned_model")
# Gradio interface
def question_answer(question):
try:
generated = text_gen(question, max_length=200, num_return_sequences=1)
generated_text = generated[0]['generated_text'].replace(question, "").strip()
return generated_text
except Exception as e:
return f"Error in generating response: {str(e)}"
text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
iface = gr.Interface(
fn=question_answer,
inputs="text",
outputs="text",
title="Answering questions about the Enron case.",
description="Ask a question about the Enron case!",
examples=["What is Eron?"]
)
iface.launch()