Wendgan/NLP_with_Disaster_Tweets

1import streamlit as st2import torch3import torch.nn as nn4import torch.nn.functional as F5import numpy as np6import re7import pickle8 9# Load word2idx10with open('word2idx.pkl', 'rb') as f:11    word2idx = pickle.load(f)12 13# Clean text function14def clean_text(text):15    emoji_pattern = re.compile("["16                               u"\U0001F600-\U0001F64F" 17                               u"\U0001F300-\U0001F5FF"  18                               u"\U0001F680-\U0001F6FF"  19                               u"\U0001F1E0-\U0001F1FF"  20                               u"\U00002702-\U000027B0"21                               u"\U000024C2-\U0001F251"22                               "]+", flags=re.UNICODE)23    text = emoji_pattern.sub(r'', text)24 25    url = re.compile(r'https?://\S+|www\.\S+')26    text = url.sub(r'', text)27 28    text = text.replace('#', ' ')29    text = text.replace('@', ' ')30 31    symbols = re.compile(r'[^A-Za-z0-9 ]')32    text = symbols.sub(r'', text)33 34    text = text.lower()35 36    return text37 38# Text to sequence function39def text_to_sequence(text, word2idx, maxlen=55):40    words = text.split()41    seq = [word2idx.get(word, 0) for word in words]42    if len(seq) > maxlen:43        seq = seq[:maxlen]44    else:45        seq = [0]*(maxlen - len(seq)) + seq46    return np.array(seq)47 48# Define the BiLSTM class49class BiLSTM(nn.Module):50    def __init__(self, weights_matrix, output_size, hidden_dim, hidden_dim2, n_layers, drop_prob=0.5):51        super(BiLSTM, self).__init__()52 53        self.output_size = output_size54        self.n_layers = n_layers55        self.hidden_dim = hidden_dim56 57        # Embedding layer58        num_embeddings, embedding_dim = weights_matrix.size()59        self.embedding = nn.Embedding(num_embeddings, embedding_dim)60        self.embedding.weight.data.copy_(weights_matrix)61        self.embedding.weight.requires_grad = False  # Freeze embedding layer62 63        # BiLSTM layer64        self.lstm = nn.LSTM(embedding_dim, hidden_dim, n_layers, dropout=drop_prob, bidirectional=True, batch_first=True)65 66        # Dropout layer67        self.dropout = nn.Dropout(0.3)68 69        # Fully connected layers70        self.fc1 = nn.Linear(hidden_dim * 2, hidden_dim2)71        self.fc2 = nn.Linear(hidden_dim2, output_size)72 73        # Activation function74        self.sigmoid = nn.Sigmoid()75 76    def forward(self, x, hidden):77        batch_size = x.size(0)78 79        # Embedding80        embeds = self.embedding(x)81 82        # LSTM83        lstm_out, hidden = self.lstm(embeds, hidden)84 85        # Stack up LSTM outputs86        lstm_out = lstm_out.contiguous().view(-1, self.hidden_dim * 2)87 88        # Dropout and fully connected layers89        out = self.dropout(lstm_out)90        out = self.fc1(out)91        out = F.relu(out)92        out = self.dropout(out)93        out = self.fc2(out)94 95        # Sigmoid activation96        sig_out = self.sigmoid(out)97 98        # Reshape to batch_size first99        sig_out = sig_out.view(batch_size, -1)100        sig_out = sig_out[:, -1]  # Get last batch of labels101 102        return sig_out, hidden103 104    def init_hidden(self, batch_size, train_on_gpu=False):105        weight = next(self.parameters()).data106 107        layers = self.n_layers * 2  # Multiply by 2 for bidirectionality108        if train_on_gpu:109            hidden = (weight.new(layers, batch_size, self.hidden_dim).zero_().cuda(),110                      weight.new(layers, batch_size, self.hidden_dim).zero_().cuda())111        else:112            hidden = (weight.new(layers, batch_size, self.hidden_dim).zero_(),113                      weight.new(layers, batch_size, self.hidden_dim).zero_())114        return hidden115 116# Load the embedding weights matrix117weights_matrix = torch.tensor(np.load('weights_matrix.npy'))118 119# Instantiate the model120output_size = 1121hidden_dim = 128122hidden_dim2 = 64123n_layers = 2124 125net = BiLSTM(weights_matrix, output_size, hidden_dim, hidden_dim2, n_layers)126 127# Load the model's state_dict128net.load_state_dict(torch.load('state_dict.pt', map_location=torch.device('cpu')))129net.eval()130 131# Streamlit app132def main():133    st.title("Disaster Tweet Classifier")134    st.write("Enter a tweet to classify whether it's about a real disaster or not.")135 136    user_input = st.text_area("Enter Tweet Text:")137 138    if st.button("Classify"):139        if user_input:140            # Preprocess input141            clean_input = clean_text(user_input)142            seq = text_to_sequence(clean_input, word2idx)143            input_tensor = torch.from_numpy(seq).unsqueeze(0).type(torch.LongTensor)144 145            # Initialize hidden state146            h = net.init_hidden(1, train_on_gpu=False)147            h = tuple([each.data for each in h])148 149            # Make prediction150            with torch.no_grad():151                output, h = net(input_tensor, h)152                prob = output.item()153                pred = int(torch.round(output).item())154 155            # Display result156            if pred == 1:157                st.success(f"This tweet is about a **real disaster**. (Probability: {prob:.4f})")158            else:159                st.info(f"This tweet is **not about a real disaster**. (Probability: {prob:.4f})")160        else:161            st.warning("Please enter some text to classify.")162 163if __name__ == '__main__':164    main()165