Upload assessment3_antonis_karaolis.py

assessment3_antonis_karaolis.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+"""Assessment3_Antonis_Karaolis.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1Qd3aOoBB6q1uy2pHPeLudMlsYd9J30-C
+"""
+
+!pip install -U sentence-transformers
+!pip install transformers
+!pip install gradio
+!pip install chromadb
+!pip install datasets
+pip install accelerate -U
+pip install transformers[torch]
+
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+import chromadb
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments
+import gradio as gr
+import torch
+from accelerate import Accelerator
+from transformers import Trainer, TrainingArguments
+from datasets import Dataset
+from torch.cuda.amp import autocast
+
+emails_df = pd.read_csv('/content/emails.csv', nrows=500, on_bad_lines='skip')
+emails_df['message'] = emails_df['message'].apply(lambda x: x.strip() if type(x) == str else '')
+
+model = SentenceTransformer('all-MiniLM-L6-v2')
+
+emails_embeddings = model.encode(emails_df['message'].tolist(), show_progress_bar=True)
+
+chroma_client = chromadb.Client()
+collection = chroma_client.create_collection(name="enron_emails_subset")
+
+collection.add(
+    embeddings=emails_embeddings.tolist(),
+    documents=emails_df['message'].tolist(),
+    metadatas=[{"email_id": idx} for idx in emails_df.index],
+    ids=[str(idx) for idx in emails_df.index]
+)
+
+tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
+tokenizer.pad_token = tokenizer.eos_token
+
+def tokenize_function(examples):
+    with autocast():
+        result = tokenizer(examples['message'], truncation=True, padding="max_length", max_length=128)
+        result["labels"] = result["input_ids"].copy()
+    return result
+
+emails_df = pd.read_csv('/content/emails.csv', nrows=500, on_bad_lines='skip')
+dataset = Dataset.from_pandas(emails_df[['message']])
+dataset = dataset.map(tokenize_function, batched=True, num_proc=4)
+
+train_dataset = dataset.train_test_split(test_size=0.1)['train']
+
+model = GPT2LMHeadModel.from_pretrained('distilgpt2')
+model.resize_token_embeddings(len(tokenizer))
+
+training_args = TrainingArguments(
+    output_dir='/content/model_output',
+    num_train_epochs=1,
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=2,
+    save_steps=250,
+    logging_dir='/content/logs',
+    logging_strategy="steps",
+    logging_steps=50
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    tokenizer=tokenizer
+)
+
+trainer.train()
+
+model.save_pretrained('/content/model_output')
+tokenizer.save_pretrained('/content/model_output')
+
+model = GPT2LMHeadModel.from_pretrained('/content/model_output')
+tokenizer = GPT2Tokenizer.from_pretrained('/content/model_output')
+
+def answer_question(question):
+    model.eval()
+    inputs = tokenizer.encode(question, return_tensors='pt')
+    outputs = model.generate(inputs, max_length=100, num_return_sequences=1)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Gradio interface
+iface = gr.Interface(fn=answer_question, inputs="text", outputs="text")
+iface.launch()