Update chat.py

chat.py

@@ -10,17 +10,22 @@ _tokenizer = None
 _model = None
 
 def _load_chatbot_resources():
+    """Loads the tokenizer and model resources only once."""
     global _tokenizer, _model
     if _tokenizer is None or _model is None:
-        _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        # FIX: Added padding_side='left' for decoder-only models (like DialoGPT) 
+        # to ensure correct token generation and suppress the warning.
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side='left')
         _model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
     return _tokenizer, _model
 
 class ChatSessionManager:
+    """Manages chat sessions, history tokens, and conversation transcripts."""
     def __init__(self) -> None:
         self._sessions: Dict[str, Dict[str, object]] = {}
 
     def _ensure_session(self, session_id: str | None) -> str:
+        """Ensures a valid session ID exists, creating a new one if necessary."""
         if not session_id or session_id not in self._sessions:
             session_id = uuid.uuid4().hex
             self._sessions[session_id] = {
@@ -30,17 +35,19 @@ class ChatSessionManager:
         return session_id
 
     def _generate_reply(self, history_tokens, message):
+        """Encodes input, generates a response using the model, and decodes the result."""
         tokenizer, model = _load_chatbot_resources()
-        # Encode user message
+        
+        # 1. Encode user message
         user_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
 
-        # Concatenate history and current input
+        # 2. Concatenate history and current input
         if history_tokens is not None:
             bot_input_ids = torch.cat([history_tokens, user_input_ids], dim=-1)
         else:
             bot_input_ids = user_input_ids
 
-        # Generate response
+        # 3. Generate response
         generated_ids = model.generate(
             bot_input_ids,
             max_length=1024,
@@ -51,13 +58,14 @@ class ChatSessionManager:
             temperature=0.8
         )
         
-        # Decode the reply
+        # 4. Decode the reply (only the new part)
         reply_ids = generated_ids[:, bot_input_ids.shape[-1]:]
         reply_text = tokenizer.decode(reply_ids[0], skip_special_tokens=True)
 
         return generated_ids, reply_text or "I am still thinking about that."
 
     def handle_message(self, session_id: str | None, message: str) -> Tuple[str, str, List[Dict[str, str]]]:
+        """Processes an incoming message, generates a reply, and updates session history."""
         session_id = self._ensure_session(session_id)
         state = self._sessions[session_id]
         transcript: List[Dict[str, str]] = state["transcript"]  # type: ignore[assignment]
@@ -77,6 +85,7 @@ class ChatSessionManager:
         return reply, session_id, list(transcript)
 
     def get_history(self, session_id: str) -> List[Dict[str, str]]:
+        """Retrieves the full transcript for a given session ID."""
         state = self._sessions.get(session_id, {"transcript": []})
         return list(state["transcript"])  # type: ignore[index]