demo / app.py
saradha12's picture
Update app.py
1b8f7c4
import imaplib
import email
from transformers import BartForConditionalGeneration, BartTokenizer, pipeline
import torch
import email.header
model_name = 'facebook/bart-large-cnn'
tokenizer = BartTokenizer.from_pretrained(model_name)
model = BartForConditionalGeneration.from_pretrained(model_name)
sentiment_analyzer = pipeline('sentiment-analysis', model='distilbert-base-uncased')
mail = imaplib.IMAP4_SSL('imap.gmail.com')
mail.login('dharsha5678@gmail.com', 'fwqw pnmq ulip umjl')
mail.select('inbox')
def generate_summary(email_text, max_length=20):
inputs = tokenizer([email_text], return_tensors='pt', max_length=1024, truncation=True)
with torch.no_grad():
summary_ids = model.generate(**inputs, max_length=max_length)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
from datetime import date
today = date.today()
today_date = today.strftime("%d-%b-%Y")
status, email_ids = mail.search(None, 'SINCE', today_date)
email_ids = email_ids[0].split()
for email_id in email_ids:
status, msg_data = mail.fetch(email_id, '(RFC822)')
raw_email = msg_data[0][1]
msg = email.message_from_bytes(raw_email)
sender = msg['From']
subject = email.header.decode_header(msg['Subject'])
subject_str = ""
for part, encoding in subject:
if isinstance(part, bytes):
if encoding:
subject_str += part.decode(encoding)
else:
subject_str += part.decode('utf-8')
else:
subject_str += part
body = ""
if msg.is_multipart():
for part in msg.walk():
if part.get_content_type() == "text/plain":
body = part.get_payload(decode=True).decode()
break
else:
body = msg.get_payload(decode=True).decode()
if body:
word_count = len(body.split())
if word_count < 10:
summary = body
else:
if word_count < 50:
summary = generate_summary(body, max_length=20)
else:
summary = generate_summary(body, max_length=50)
sentiment_result = sentiment_analyzer(summary)
label = sentiment_result[0]['label']
score = sentiment_result[0]['score']
if score >= 0.53:
email_label = "Important"
else:
email_label = "Not Important"
print(f"From: {sender}")
print(f"Email Subject: {subject_str}")
print(f"Generated Summary: {summary}")
print(f"Sentiment Label: {email_label}")
print(f"Sentiment Score: {score}")
print("-" * 50)
mail.logout()