Update app.py

app.py

@@ -7,34 +7,51 @@ 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)
+train_content, _ = train_test_split(content_text, train_size=0.00005) # was unable to load more emails
 
 # 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
+# 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()
 
-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)
+# saving the model
+model.save_pretrained("./fine_tuned_model")
+tokenizer.save_pretrained("./fine_tuned_model")
 
+# Gradio interface
 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()
@@ -42,12 +59,13 @@ def question_answer(question):
     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 Enron?"]
+    examples=["What is Eron?"]
 )
 iface.launch()