|
from transformers import T5Tokenizer, T5ForConditionalGeneration |
|
|
|
from transformers import AdamW |
|
import pandas as pd |
|
import torch |
|
import pytorch_lightning as pl |
|
from pytorch_lightning.callbacks import ModelCheckpoint |
|
from torch.nn.utils.rnn import pad_sequence |
|
|
|
|
|
pl.seed_everything(100) |
|
|
|
MODEL_NAME='t5-base' |
|
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
INPUT_MAX_LEN = 128 |
|
OUTPUT_MAX_LEN = 128 |
|
|
|
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, model_max_length=512) |
|
|
|
class T5Model(pl.LightningModule): |
|
|
|
def __init__(self): |
|
super().__init__() |
|
self.model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict = True) |
|
|
|
|
|
def forward(self, input_ids, attention_mask, labels=None): |
|
|
|
output = self.model( |
|
input_ids=input_ids, |
|
attention_mask=attention_mask, |
|
labels=labels |
|
) |
|
return output.loss, output.logits |
|
|
|
def training_step(self, batch, batch_idx): |
|
|
|
input_ids = batch["input_ids"] |
|
attention_mask = batch["attention_mask"] |
|
labels= batch["target"] |
|
loss, logits = self(input_ids , attention_mask, labels) |
|
|
|
|
|
self.log("train_loss", loss, prog_bar=True, logger=True) |
|
|
|
return {'loss': loss} |
|
|
|
def validation_step(self, batch, batch_idx): |
|
input_ids = batch["input_ids"] |
|
attention_mask = batch["attention_mask"] |
|
labels= batch["target"] |
|
loss, logits = self(input_ids, attention_mask, labels) |
|
|
|
self.log("val_loss", loss, prog_bar=True, logger=True) |
|
|
|
return {'val_loss': loss} |
|
|
|
def configure_optimizers(self): |
|
return AdamW(self.parameters(), lr=0.0001) |
|
|
|
train_model = T5Model.load_from_checkpoint('best-model.ckpt',map_location=DEVICE) |
|
train_model.freeze() |
|
|
|
def generate_response(question): |
|
|
|
inputs_encoding = tokenizer( |
|
question, |
|
add_special_tokens=True, |
|
max_length= INPUT_MAX_LEN, |
|
padding = 'max_length', |
|
truncation='only_first', |
|
return_attention_mask=True, |
|
return_tensors="pt" |
|
) |
|
|
|
|
|
generate_ids = train_model.model.generate( |
|
input_ids = inputs_encoding["input_ids"], |
|
attention_mask = inputs_encoding["attention_mask"], |
|
max_length = INPUT_MAX_LEN, |
|
num_beams = 4, |
|
num_return_sequences = 1, |
|
no_repeat_ngram_size=2, |
|
early_stopping=True, |
|
) |
|
|
|
preds = [ |
|
tokenizer.decode(gen_id, |
|
skip_special_tokens=True, |
|
clean_up_tokenization_spaces=True) |
|
for gen_id in generate_ids |
|
] |
|
|
|
return "".join(preds) |
|
|
|
import uuid |
|
import datetime |
|
import os |
|
import streamlit as st |
|
from streamlit_chat import message |
|
|
|
|
|
from pymongo.mongo_client import MongoClient |
|
from pymongo.server_api import ServerApi |
|
|
|
password=os.getenv("mongo_pass") |
|
uri = "mongodb+srv://rohank587:"+password+"@rkcluster.e3fpzja.mongodb.net/?retryWrites=true&w=majority" |
|
|
|
|
|
client = MongoClient(uri, server_api=ServerApi('1')) |
|
|
|
st.title(":red[_Sarcastic_] Chatbot") |
|
|
|
if 'generated' not in st.session_state: |
|
st.session_state['generated'] = [] |
|
if 'past' not in st.session_state: |
|
st.session_state['past'] = [] |
|
if 'messages' not in st.session_state: |
|
st.session_state['messages'] = [ |
|
{"role": "system", "content": "You are a helpful assistant."} |
|
] |
|
|
|
|
|
|
|
|
|
response_container = st.container() |
|
|
|
container = st.container() |
|
|
|
with container: |
|
with st.form(key='my_form', clear_on_submit=True): |
|
user_input = st.text_input("You:", key='input',placeholder="Disclaimer: Be careful with punctuations like , ? . ! \"") |
|
submit_button = st.form_submit_button(label='Send',use_container_width=True) |
|
col1,col2=st.columns(2) |
|
with col1: |
|
clear_button = st.button("Clear Conversation", key="clear",use_container_width=True) |
|
with col2: |
|
save_button = st.button("Save Conversation", key="save",use_container_width=True) |
|
down_id = st.text_input('Enter ID to download chat',placeholder="Message ID") |
|
if down_id: |
|
info=client['rohank']['table1'] |
|
data=info.find_one({'message_id':down_id}) |
|
down_button = st.download_button('Download chat', "\n".join(data['message']),file_name="sar_chat.txt") |
|
|
|
if clear_button: |
|
st.session_state['generated'] = [] |
|
st.session_state['past'] = [] |
|
st.session_state['messages'] = [ |
|
{"role": "system", "content": "You are a helpful assistant."} |
|
] |
|
|
|
if save_button and st.session_state['generated'] and st.session_state['past']: |
|
|
|
try: |
|
client.admin.command('ping') |
|
st.success("Pinged your deployment. You successfully connected to MongoDB! Saved Successfully.") |
|
info=client['rohank']['table1'] |
|
chats=list([]) |
|
for i in range(len(st.session_state['generated'])): |
|
chats.append("You: "+st.session_state['past'][i]) |
|
chats.append("Bot: "+st.session_state['generated'][i]) |
|
id=uuid.uuid4() |
|
time=datetime.datetime.now() |
|
info.insert_one({"time of saving":time.strftime("%c"),"message_id":str(id),"message":chats}) |
|
st.success("Copy this id "+str(id)+" for downloading saved chat anytime anywhere and then paste it down below!") |
|
except Exception as e: |
|
st.error("Can't connect to MongoDB. Save Failed.") |
|
|
|
|
|
if submit_button and user_input: |
|
output = generate_response(user_input) |
|
st.session_state['past'].append(user_input) |
|
st.session_state['generated'].append(output) |
|
|
|
if st.session_state['generated']: |
|
with response_container: |
|
for i in range(len(st.session_state['generated'])): |
|
message(st.session_state["past"][i], is_user=True, key=str(i) + '_user') |
|
message(st.session_state["generated"][i], key=str(i)) |