Update app.py

app.py

@@ -1,104 +1,18 @@
 # app.py
 
-import streamlit as st
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import json
-import pickle
-
-# ─── Model definitions ─────────────────────────────────────────────────────────
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, max_len=32):
-        super().__init__()
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model))
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        self.pe = pe.unsqueeze(0)
-    def forward(self, x):
-        return x + self.pe[:, : x.size(1)].to(x.device)
-
-class EmotionTransformer(nn.Module):
-    def __init__(self, vocab_size, embed_dim, num_heads, num_classes, pad_idx, max_len=32):
-        super().__init__()
-        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
-        self.pos_enc = PositionalEncoding(embed_dim, max_len)
-        layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, batch_first=True)
-        self.trans_enc = nn.TransformerEncoder(layer, num_layers=2)
-        self.dropout = nn.Dropout(0.3)
-        self.fc = nn.Linear(embed_dim, num_classes)
-    def forward(self, x):
-        mask = x == pad_idx
-        x = self.embedding(x)
-        x = self.pos_enc(x)
-        x = self.trans_enc(x, src_key_padding_mask=mask)
-        x = self.dropout(x.mean(dim=1))
-        return self.fc(x)
-
-# ─── Load artifacts ────────────────────────────────────────────────────────────
-
-@st.cache(allow_output_mutation=True)
-def load_model_and_artifacts():
-    # load vocab and label encoder
-    with open("vocab.json", "r") as vf:
-        vocab = json.load(vf)
-    with open("label_encoder.pkl", "rb") as lf:
-        le = pickle.load(lf)
-    # instantiate & load model
-    pad_idx = vocab.get("<PAD>", 0)
-    num_classes = len(le.classes_)
-    model = EmotionTransformer(
-        vocab_size=len(vocab),
-        embed_dim=64,
-        num_heads=4,
-        num_classes=num_classes,
-        pad_idx=pad_idx,
-        max_len=32,
-    )
-    state = torch.load("emotion_transformer_model.pth", map_location="cpu")
-    model.load_state_dict(state)
-    model.eval()
-    return model, vocab, le
-
-model, vocab, le = load_model_and_artifacts()
-
-# ─── Preprocessing & Inference ─────────────────────────────────────────────────
-
-def predict_emotion(text: str):
-    # simple clean + tokenize
-    tokens = text.lower().split()
-    ids = [vocab.get(w, vocab.get("<UNK>", 1)) for w in tokens]
-    # pad/truncate
-    if len(ids) < 32:
-        ids = ids + [vocab.get("<PAD>", 0)] * (32 - len(ids))
-    else:
-        ids = ids[:32]
-    x = torch.tensor([ids])
-    with torch.no_grad():
-        logits = model(x)
-        probs = F.softmax(logits, dim=1)[0]
-        idx = torch.argmax(probs).item()
-    return le.inverse_transform([idx])[0]
-
-# ─── UI ────────────────────────────────────────────────────────────────────────
-
-st.set_page_config(page_title="EmotiBot Lite", layout="centered")
-st.title("🌿 EmotiBot (Lite Inference)")
-
-user_input = st.text_input("How are you feeling?")
-
-if user_input:
-    emotion = predict_emotion(user_input)
-    st.write(f"**Predicted emotion:** {emotion}")
-
-    suggestions = {
-        "sadness": "It’s okay to feel down sometimes. I’m here to support you.",
-        "anger": "That must have been frustrating. Want to vent about it?",
-        "love": "That’s beautiful to hear! What made you feel that way?",
-        "happiness": "That's awesome! What’s bringing you joy today?",
-        "neutral": "Got it. I’m here if you want to share more.",
-    }
-    st.write("**Advice:**", suggestions.get(emotion, "I’m here for you."))
+from src.inference import predict
+from src.responses import get_response
+
+print("EmotiBot 🌿: Hi! How are you feeling today? (Type 'exit' to quit)")
+
+while True:
+    user_raw = input("You: ").strip()
+    if user_raw.lower() in ["exit", "quit"]:
+        print("EmotiBot 🌿: Take care! I’m here whenever you want to talk.")
+        break
+
+    emotion = predict(user_raw)
+    reply, done = get_response(emotion, user_raw)
+    print(f"EmotiBot 🌿: {reply}")
+    if done:
+        break