Update app.py

app.py

@@ -1,155 +1,96 @@
-import os
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoConfig
-import torch.nn as nn
-
-# ----- Model Definition -----
-class CustomDialoGPT(nn.Module):
-    def __init__(self, vocab_size, n_embd=768, n_head=12, n_layer=12):
-        super().__init__()
-
-        config = AutoConfig.from_pretrained("microsoft/DialoGPT-medium",
-            vocab_size=vocab_size,
-            n_embd=n_embd,
-            n_head=n_head,
-            n_layer=n_layer,
-            bos_token_id=50256,
-            eos_token_id=50256,
-            pad_token_id = 50256
-        )
-        self.transformer = AutoModelForCausalLM.from_config(config)
-        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
-
-    def forward(self, input_ids):
-        transformer_outputs = self.transformer(input_ids=input_ids, output_hidden_states=True)
-        hidden_states = transformer_outputs.hidden_states[-1] #get last hidden state
-        logits = self.lm_head(hidden_states)
-        return logits
-    
-    def get_num_params(self):
-        return sum(p.numel() for p in self.parameters())
-    
-def build_model(vocab_size, target_params=128_000_000):
-    """Build a model with around the target parameter count, ensuring n_embd is divisible by n_head."""
-    
-    n_embd_options = [512, 768, 1024]
-    n_head_options = [8, 12, 16]
-    n_layer_options = [6, 8, 12, 16]
-
-    best_params_diff = float('inf')
-    best_n_embd = 0
-    best_n_head = 0
-    best_n_layer = 0
-    
-    for n_embd in n_embd_options:
-      for n_head in n_head_options:
-          if n_embd % n_head != 0:
-              continue
-          for n_layer in n_layer_options:
-              model = CustomDialoGPT(vocab_size, n_embd, n_head, n_layer)
-              params = model.get_num_params()
-              params_diff = abs(params - target_params)
-              
-              if params_diff < best_params_diff:
-                  best_params_diff = params_diff
-                  best_n_embd = n_embd
-                  best_n_head = n_head
-                  best_n_layer = n_layer
-              
-              del model
-    print("Model parameters:", best_n_embd, best_n_head, best_n_layer)
-
-    return CustomDialoGPT(vocab_size, best_n_embd, best_n_head, best_n_layer)
-
-def chat_with_model(user_input, model, tokenizer, device="cpu", max_length=100, temperature=0.8):
-    """
-    Generates a response from the model given the user input.
-    """
-    input_ids = tokenizer.encode(user_input, return_tensors='pt').to(device)
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+from huggingface_hub import hf_hub_download
+import os
+
+# Model and tokenizer details
+model_repo = "elapt1c/ElapticAI-1a"
+model_filename = "pytorch_model.bin" # Assuming the model is saved as pytorch_model.bin, adjust if needed. Check the HF repo.
+tokenizer_name = "microsoft/DialoGPT-medium"
+
+# Device configuration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+# Load model configuration
+config = AutoConfig.from_pretrained("microsoft/DialoGPT-medium")
+# Initialize model from config (important to use the same architecture)
+model = AutoModelForCausalLM.from_config(config)
+
+# Download and load model weights
+try:
+    pth_filepath = hf_hub_download(repo_id=model_repo, filename=model_filename)
+    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}")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    print("Please ensure the model repository and filename are correct.")
+    raise e # It's better to raise the error in a Space, so it's visible.
+
+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)
 
     with torch.no_grad():
-        output = model.transformer.generate(
-            inputs=input_ids,
-            max_length=max_length,
-            temperature=temperature,
+        output = model.generate(
+            input_ids,
+            max_length=1000, # Adjust as needed
             pad_token_id=tokenizer.eos_token_id,
+            temperature=0.7,
+            top_p=0.9
         )
-        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return generated_text
-        
-def load_model_and_tokenizer(model_repo, device):
-    """Loads the model and tokenizer from the Hugging Face model repo."""
-    try:
-        # Check if running in hugging face
-        if 'HF_MODEL_ID' in os.environ:
-            # Load tokenizer
-            tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-            vocab_size = len(tokenizer)
-            
-            model = build_model(vocab_size)
-
-            #Load model and optimizer
-            checkpoint_files = [f for f in os.listdir(".") if f.endswith('.pth')]
-            if not checkpoint_files:
-                print("No checkpoint found. Please train the model first.")
-                return
-            
-            checkpoint_path = checkpoint_files[0]
-            
-            checkpoint = torch.load(checkpoint_path, map_location=device)
-            model.load_state_dict(checkpoint['model_state_dict'])
-
-            model.to(device)
-            model.eval()
-            print(f"Model loaded on device: {device}")
-        else:
-            # Load tokenizer
-            tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-            vocab_size = len(tokenizer)
-            
-            model = build_model(vocab_size)
-            
-            #Load model and optimizer
-            checkpoint_path = input("Enter the path to your .pth checkpoint file: ")
-            
-            checkpoint = torch.load(checkpoint_path, map_location=device)
-            model.load_state_dict(checkpoint['model_state_dict'])
-
-            model.to(device)
-            model.eval()
-            print(f"Model loaded on device: {device}")
-        return model, tokenizer
-    except Exception as e:
-        print(f"Error loading model or tokenizer: {e}")
-        return None, None
-
-
-def gradio_chat(model, tokenizer, device="cpu", max_length = 100, temperature = 0.8):
-  """Defines the gradio chatbot interaction."""
-  def respond(message, chat_history):
-    bot_message = chat_with_model(message, model, tokenizer, device=device, max_length=max_length, temperature = temperature)
-    chat_history.append((message, bot_message))
-    return "", chat_history
-
-  with gr.Blocks() as demo:
-        chatbot = gr.Chatbot()
-        msg = gr.Textbox()
-        clear = gr.Button("Clear")
-        
-        msg.submit(respond, [msg, chatbot], [msg, chatbot])
-        clear.click(lambda: None, None, chatbot, queue=False)
-
-  return demo
 
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # 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()
+
+    history.append((user_input, bot_response))
+    return bot_response, history
+
+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
+
+
+iface = gr.ChatInterface(
+    fn=chat_with_model,
+    inputs=gr.Chatbox(placeholder="Type your message here..."),
+    outputs=gr.Chatbot(),
+    title="ElapticAI-1a Chatbot",
+    description="Simple chatbot interface for ElapticAI-1a model. Talk to the model and see its responses!",
+    examples=[
+        ["Hello"],
+        ["How are you?"],
+        ["Tell me a joke"]
+    ]
+)
 
 if __name__ == "__main__":
-    # Load model and tokenizer (downloading from hugging face model repo).
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-    
-    model, tokenizer = load_model_and_tokenizer("elapt1c/ElapticAI-1a", device = device)
-    if model and tokenizer:
-      #launch the gradio interface.
-      demo = gradio_chat(model, tokenizer, device = device)
-      demo.launch()
+    iface.launch()