Update app.py

app.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 
 # ----- Model Definition -----
 class CustomDialoGPT(nn.Module):
-    def __init__(self, vocab_size, n_embd=768, n_head=8, n_layer=8): # <---- FORCE n_embd, n_head, n_layer to match DialoGPT-medium
+    def __init__(self, vocab_size, n_embd=768, n_head=8, n_layer=8): # <---- FORCE n_embd, n_head, n_layer to match your model
         super().__init__()
 
         config = AutoConfig.from_pretrained("microsoft/DialoGPT-medium",
@@ -39,14 +39,13 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-vocab_size = len(tokenizer)
+vocab_size = len(tokenizer) # <---- Define vocab_size AFTER loading tokenizer
 
 # Initialize model with fixed parameters to match checkpoint
-n_embd=768 # <---- FORCE n_embd to 768
-n_head=8  # <---- FORCE n_head to 12
-n_layer=8 # <---- FORCE n_layer to 12
-model = CustomDialoGPT(vocab_size, n_embd, n_head, n_layer)
-
+n_embd=768
+n_head=8
+n_layer=8
+model = CustomDialoGPT(vocab_size, n_embd, n_head, n_layer).to(device).eval()
 
 # Download and load model weights
 try:
@@ -54,13 +53,11 @@ try:
     checkpoint = torch.load(pth_filepath, map_location=device)
 
     # Handle different checkpoint saving formats if needed.
-    # If your checkpoint is just the state_dict, load it directly.
     if 'model_state_dict' in checkpoint:
         model.load_state_dict(checkpoint['model_state_dict'])
     elif 'state_dict' in checkpoint:
         model.load_state_dict(checkpoint['state_dict'])
     else:
-        # Assume checkpoint is just the raw state_dict
         model.load_state_dict(checkpoint)
 
     print(f"Successfully loaded model weights from {model_repo}/{model_filename}")
@@ -72,60 +69,38 @@ except Exception as e:
 model.to(device)
 model.eval() # Set model to evaluation mode
 
-def chat_with_model(user_input, history=[]):
-    """Chatbot function to interact with the loaded model."""
-    history_transformer_format = history_to_transformer_format(history)
-    input_text = tokenizer.eos_token.join(history_transformer_format + [user_input])
-
-    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)
+def chat_with_model(user_input): # Removed history parameter for gr.Text() output
+    """Chatbot function to interact with the loaded model - DYNAMIC RESPONSE."""
+    input_ids = tokenizer.encode(user_input, return_tensors='pt').to(device)
 
     with torch.no_grad():
-        output = model.transformer.generate( # Use model.transformer.generate here
-            inputs=input_ids, # Use inputs instead of input_ids
-            max_length=1000, # Adjust as needed
+        output = model.transformer.generate(
+            inputs=input_ids,
+            max_length=100,
             pad_token_id=tokenizer.eos_token_id,
             temperature=0.7,
-            top_p=0.9
+            top_p=0.9,
+            do_sample=True
         )
 
     response = tokenizer.decode(output[0], skip_special_tokens=True)
+    bot_response = response # No need to split for gr.Text()
 
-    # Extract only the bot's last response, assuming it's after the last user input.
-    # This is a simple heuristic and might need adjustments based on training data format.
-    split_response = response.split(tokenizer.eos_token)
-    bot_response = split_response[-1].strip()
-
-    # Explicitly format history as list of tuples:
-    history.append((user_input, bot_response))
-    
-    # Reformat history for Gradio Chatbot - Ensure tuples within a list
-    chatbot_history = []
-    for turn in history:
-        chatbot_history.append(turn) # Each turn is already a tuple (user_msg, bot_msg)
-
-    return bot_response, chatbot_history # Return chatbot_history for Gradio
+    print("--- chat_with_model Output ---") # Debugging print
+    print("user_input:", user_input)        # Debugging print
+    print("bot_response:", bot_response)      # Debugging print
+    print("--- End of chat_with_model Output ---") # Debugging print
 
-def history_to_transformer_format(history):
-    """Convert gradio history to a list of strings for transformer input."""
-    history_formatted = []
-    for user_msg, bot_msg in history:
-        history_formatted.append(user_msg)
-        history_formatted.append(bot_msg)
-    return history_formatted
+    return bot_response # Just return bot_response for gr.Text()
 
 
 iface = gr.Interface( # Changed from gr.ChatInterface to gr.Interface
     fn=chat_with_model,
     inputs=gr.Textbox(placeholder="Type your message here..."), # Explicitly define inputs as gr.Textbox
-    outputs=gr.Chatbot(), # Explicitly define outputs as gr.Chatbot
-    title="ElapticAI-1a Chatbot",
-    description="Simple chatbot interface for ElapticAI-1a model. Talk to the model and see its responses!",
-    examples=[ # Corrected examples format
-        ["Hello", "Hi there!"], # Example 1: [user_input, bot_response]
-        ["How are you?", "I am doing well, thank you."], # Example 2
-        ["Tell me a joke", "Why don't scientists trust atoms? Because they make up everything! 😄"] # Example 3
-    ]
+    outputs=gr.Text(), # Changed outputs to gr.Text()
+    title="ElapticAI-1a Chatbot - TESTING MODEL RESPONSE", # Updated title
+    description="Simple chatbot interface for ElapticAI-1a model - TESTING MODEL RESPONSE" # Updated description
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()