Update app.py

app.py

@@ -2,36 +2,60 @@ import gradio as gr
 import torch
 import tiktoken
 
-from model import GPTLanguageModel
+from model import GPTLanguageModel  # Import the model from model.py
 
 # Load the model and tokenizer
 def load_model():
     """Load the trained GPT model"""
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
     GPT_CONFIG = {
-        "vocab_size" : 50257,
-        "n_heads" : 8,
-        "n_layers" : 6,
-        "head_size" : 64,
-        "n_embd" : 512,
-        "block_size" : 128,
-        "dropout" : 0.1,
-        "learning_rate" : 3e-4,
-        "weight_decay" : 0.1,
+        "vocab_size": 50257,
+        "n_heads": 8,
+        "n_layers": 6,
+        "n_embd": 512,
+        "block_size": 128,
+        "dropout": 0.1,
     }
-    model = GPTLanguageModel(GPT_CONFIG)
+
+    model = GPTLanguageModel(
+        GPT_CONFIG["vocab_size"],
+        GPT_CONFIG["n_embd"],
+        GPT_CONFIG["block_size"],
+        GPT_CONFIG["n_layers"],
+        GPT_CONFIG["n_heads"],
+        device
+    )
+    
+    # Load the trained weights
     model.load_state_dict(torch.load("model_weights.pth", map_location=device))
     model.to(device)
-    model.eval()
+    model.eval()  # Set the model to evaluation mode
+
+    # Use tiktoken for tokenization
     tokenizer = tiktoken.get_encoding("gpt2")
+
     return model, tokenizer, device
 
-# Load model globally
+# Load the model globally
 model, tokenizer, device = load_model()
 
-# Define the respond function
+# Tokenization and detokenization functions
+def text_to_token_ids(text, tokenizer):
+    return torch.tensor([tokenizer.encode(text)], dtype=torch.long)
+
+def token_ids_to_text(token_ids, tokenizer):
+    return tokenizer.decode(token_ids[0].tolist())
+
+# Generate text function using the model
+def generate_text(model, idx, max_new_tokens, context_size=256):
+    # Call the model's generate function
+    token_ids = model.generate(idx, max_new_tokens)
+    return token_ids
+
+# Define the response function
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens):
-    # Build message history with system message
+    # Build the message history with the system message
     messages = [{"role": "system", "content": system_message}]
     
     for val in history:
@@ -43,8 +67,11 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens)
     # Add the user message to the conversation
     messages.append({"role": "user", "content": message})
 
+    # Concatenate the history into one context
+    conversation_history = " ".join([msg["content"] for msg in messages])
+
     # Convert the latest user message to token IDs
-    input_ids = text_to_token_ids(message, tokenizer).to(device)
+    input_ids = text_to_token_ids(conversation_history, tokenizer).to(device)
 
     # Generate the response from the model
     token_ids = generate_text(
@@ -56,16 +83,17 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens)
 
     # Convert the token IDs back to text and return
     response_text = token_ids_to_text(token_ids, tokenizer)
+    
     return response_text
 
-# Gradio ChatInterface
+# Gradio Chat Interface
 demo = gr.ChatInterface(
-    respond,
+    fn=respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=256, value=50, step=1, label="Max new tokens")
+        gr.Textbox(value="You are a friendly chatbot.", label="System message"),  # System message input
+        gr.Slider(minimum=1, maximum=256, value=50, step=1, label="Max new tokens")  # Max tokens slider
     ]
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)