Spaces:
Runtime error
Runtime error
File size: 2,308 Bytes
b15e74b |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import pandas as pd
import numpy as np
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModel
import torch
from sentence_transformers import SentenceTransformer, util
from tqdm import tqdm
from peewee import SqliteDatabase, Model, TextField
# Load the Enron Email Dataset
emails_df = pd.read_csv("/content/emails.csv")
# Define the ChromaDB database
db = SqliteDatabase('email_embeddings.db')
# Define the model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
# Fine-tune the Language Model on the dataset
# Fine-tune the Language Model on the dataset
# Tokenize the dataset
tokenized_texts = [tokenizer.encode(text, return_tensors="pt", max_length=512) for text in emails_df['Message']]
# Convert tokenized texts to tensors
input_ids = torch.cat(tokenized_texts, dim=1)
# Define a PyTorch dataset
dataset = torch.utils.data.TensorDataset(input_ids)
# Define the Sentence Transformer model
sentence_model = SentenceTransformer('distilbert-base-nli-mean-tokens')
# Function to create embeddings of the email dataset and store them in the ChromaDB database
def create_embeddings():
db.connect()
db.create_tables([Email])
embeddings = []
for index, row in tqdm(emails_df.iterrows(), total=len(emails_df)):
text = row['Message']
embeddings.append(sentence_model.encode(text))
for index, embedding in enumerate(embeddings):
Email.create(id=index, embedding=embedding.tobytes())
db.close()
# Define the Gradio Interface
def answer_question(question):
# Encode the question
inputs = tokenizer(question, return_tensors="pt", max_length=512, truncation=True)
# Generate response using the model
outputs = model.generate(**inputs)
# Decode the response
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# Define the Peewee Model for the ChromaDB database
class BaseModel(Model):
class Meta:
database = db
class Email(BaseModel):
embedding = TextField()
# Create a Gradio Interface
gr.Interface(fn=answer_question, inputs="text", outputs="text").launch()
# Uncomment the line below to create embeddings and store them in the ChromaDB database
create_embeddings()
|