Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,102 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import numpy as np
+import tensorflow as tf
+from transformers import RobertaTokenizerFast, TFRobertaModel
+
+MAX_LEN = 96
+MODEL_PATH = "src/roberta_fold_1.h5" 
+
+#MODEL ARCHITECTURE
+def build_model():
+    ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name="input_ids")
+    att = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name="attention_mask")
+
+    roberta = TFRobertaModel.from_pretrained('roberta-base')
+    x = roberta(ids, attention_mask=att)[0]
+    x1 = tf.keras.layers.Dropout(0.1)(x) 
+    start_logits = tf.keras.layers.Dense(1)(x1)
+    start_logits = tf.keras.layers.Flatten()(start_logits)
+    end_logits = tf.keras.layers.Dense(1)(x1)
+    end_logits = tf.keras.layers.Flatten()(end_logits)
+    
+    model = tf.keras.Model(inputs=[ids, att], outputs=[start_logits, end_logits])
+    return model
+
+#MODEL
+@st.cache_resource
+def load_model_and_tokenizer():
+    tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base')
+    model = build_model()
+    model.load_weights(MODEL_PATH)
+    return model, tokenizer
+
+with st.spinner('Loading model and tokenizer... Please wait.'):
+    model, tokenizer = load_model_and_tokenizer()
+
+st.success("Model Loaded Successfully! 🚀")
+
+#PREDICTION 
+def predict_sentiment(tweet, sentiment):
+    # Preprocessing
+    tweet = " " + " ".join(str(tweet).split())
+    # Tokenization
+    enc = tokenizer.encode_plus(
+        sentiment, 
+        tweet,
+        add_special_tokens=True,
+        max_length=MAX_LEN,
+        padding='max_length',
+        truncation=True,
+        return_attention_mask=True)
+    
+    input_ids = np.array([enc['input_ids']])
+    attention_mask = np.array([enc['attention_mask']])
+    # Inference
+    preds = model.predict([input_ids, attention_mask])
+    start_logits = preds[0]
+    end_logits = preds[1]
+    
+    idx_start = np.argmax(start_logits)
+    idx_end = np.argmax(end_logits)
+    # Decoding
+    if sentiment == "neutral" or len(tweet.split()) < 2:
+        return tweet.strip()
+    # Validation check
+    if idx_start > idx_end:
+        return tweet.strip()
+    else:
+        text_tokens = tokenizer.decode(enc['input_ids'][idx_start:idx_end+1])
+        return text_tokens.strip()
+
+#INTERFACE
+st.title("🐦 Tweet Sentiment Extraction")
+st.markdown("""
+This AI model analyzes a tweet and extracts the **specific phrase or word** that supports the selected sentiment.
+*Built with TensorFlow & RoBERTa*
+""")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+#Inputs
+col1, col2 = st.columns([1, 2])
+with col1:
+    sentiment = st.selectbox(
+        "Select Sentiment:", 
+        ["positive", "negative", "neutral"])
+with col2:
+    text = st.text_area(
+        "Enter Tweet:", 
+        value="This movie was really fantastic and I loved it!",
+        height=100)
+#Prediction
+if st.button("Extract Phrase", type="primary"):
+    if text:
+        result = predict_sentiment(text, sentiment)
+        st.write("### Result:")
+        # Highlight Logic for Visualization
+        if result in text and len(result) > 2:
+            highlighted_text = text.replace(result, f"<span style='background-color: #ffd700; color: black; padding: 2px; font-weight: bold; border-radius: 3px;'>{result}</span>")
+            st.markdown(f"**Context:** {highlighted_text}", unsafe_allow_html=True)
+        else:
+            st.markdown(f"**Context:** {text}")  
+        st.success(f"Extracted Text: **{result}**")
+    else:
+        st.warning("Please enter a tweet to analyze.")