Update app.py

app.py

@@ -15,7 +15,7 @@ nltk.download(['punkt', 'punkt_tab'], quiet=True)
 DEVICE = torch.device("cpu")
 
 CACHE_FILE = "ubuntu_data_cache.pt"       # To get the Vocab from cache
-MODEL_FILE = "ubuntu_chatbot_best.pt"     # trained model 
+MODEL_FILE_WITH_ATTN = "ubuntu_chatbot_with_attn.pt"     # trained model with attn
 
 
 # ------------- tokenization + helpers -------------
@@ -110,7 +110,7 @@ class Encoder(nn.Module):
         return out, h                    
 
 
-class Decoder(nn.Module):
+class Decoder_with_attn(nn.Module):
     def __init__(self):
         super().__init__()
         self.emb = nn.Embedding(len(vocab), 256, padding_idx=PAD_IDX)
@@ -147,12 +147,28 @@ class Decoder(nn.Module):
         logits = self.out(out)                   
         return logits, hidden
 
-
-class Model(nn.Module):
+class Decoder_no_attn(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.emb = nn.Embedding(len(vocab), 256, padding_idx=0)
+        self.dropout = nn.Dropout(0.3) # added dropout layer
+        self.gru = nn.GRU(256, 512, num_layers=2, batch_first=True)
+     
+        self.out = nn.Linear(512, len(vocab))
+        self.norm = nn.LayerNorm(512)
+        
+    def forward(self, inp, hidden):
+        e = self.dropout(self.emb(inp))
+ 
+        out, hidden = self.gru(e, hidden)
+        out = self.norm(out.squeeze(1))
+        return self.out(out), hidden
+
+class Model_with_attn(nn.Module):
     def __init__(self):
         super().__init__()
         self.encoder = Encoder()
-        self.decoder = Decoder()
+        self.decoder = Decoder_with_attn()
 
     def forward(self, src, tgt, tf=0.5):
         enc_out, h = self.encoder(src)
@@ -168,21 +184,21 @@ class Model(nn.Module):
 
 
 # ------------- load trained model -------------
-if not os.path.exists(MODEL_FILE):
+if not os.path.exists(MODEL_FILE_WITH_ATTN):
     raise FileNotFoundError(
-        f"{MODEL_FILE} not found in Space. Upload your ubuntu_chatbot_best.pt checkpoint."
+        f"{MODEL_FILE_WITH_ATTN} not found in Space. Upload your ubuntu_chatbot_best.pt checkpoint."
     )
 
-model = Model().to(DEVICE)
-ckpt = torch.load(MODEL_FILE, map_location="cpu")
-model.load_state_dict(ckpt["model"])
-model.eval()
+model_with_attn = Model_with_attn().to(DEVICE)
+ckpt = torch.load(MODEL_FILE_WITH_ATTN, map_location="cpu")
+model_with_attn.load_state_dict(ckpt["model"])
+model_with_attn.eval()
 
 print("Model and vocab loaded. Chatbot ready to serve ")
 
 
 # ------------- beam search (beam_generate_v2 from notebook) -------------
-def beam_generate_v2(src_tensor, beam=5, max_len=50, alpha=0.7):
+def beam_generate_v2(model,src_tensor, beam=5, max_len=50, alpha=0.7):
     """
     src_tensor: [1, T] LongTensor with <SOS> ... <EOS>
     alpha: length penalty factor
@@ -258,18 +274,18 @@ def beam_generate_v2(src_tensor, beam=5, max_len=50, alpha=0.7):
 
 
 # ------------- wrapper to go from user text → reply -------------
-def generate_reply_no_attn(user_text: str) -> str:
+def generate_reply_attn(user_text: str) -> str:
     # replicate notebook logic: reverse the input sentence
     user_text_rev = reverse(user_text)
     tokens = tokenize(user_text_rev)
     ids = [SOS_IDX] + [vocab.word2idx.get(w, UNK_IDX) for w in tokens] + [EOS_IDX]
     src = torch.tensor([ids], dtype=torch.long, device=DEVICE)
-    reply = beam_generate_v2(src, beam=5, max_len=50)
+    reply = beam_generate_v2(model_with_attn,src, beam=5, max_len=50)
     if not reply.strip():
         return "I'm a chatbot trained on Ubuntu Linux support conversations, so I may not understand this question."
     return reply
 
-def generate_reply_attn(user_text: str) -> str:
+def generate_reply_no_attn(user_text: str) -> str:
     """
     Inference using the ATTENTION model.
     Replace body with your encoder/decoder calls (beam or greedy).