File size: 1,984 Bytes
746b99c c90606a 746b99c c90606a 746b99c c90606a 746b99c c90606a 746b99c c90606a 746b99c c90606a 746b99c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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)
# 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))])
# initialize model and tokenizer globally but don't load them yet
tokenizer = None
model = None
text_gen = None
def load_model():
global tokenizer, model, text_gen
if model is None or tokenizer is None:
tokenizer = GPT2Tokenizer.from_pretrained('./fine_tuned_model')
model = GPT2LMHeadModel.from_pretrained('./fine_tuned_model')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
def question_answer(question):
load_model() # loading model on first use
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)}"
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 Enron?"]
)
iface.launch() |