Update app.py

app.py

@@ -1,15 +1,19 @@
 from typing import List, Tuple, Dict, Generator
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2TokenizerFast
 import torch
 import gradio as gr
 
-# Load your safetensors model and tokenizer
-model_name = "DuckyPolice/ElapticAI-1a"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+# Load the GPT-2 tokenizer
+tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
 
+# Load the saved GPT-2 model from the local checkpoint
+model_path = "DuckyPolice/ElapticAI-1a"  # Adjust to your specific model path if needed
+model = GPT2LMHeadModel.from_pretrained(model_path)
+
+# Move model to appropriate device (GPU if available, otherwise CPU)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
+model.eval()
 
 def create_history_messages(history: List[Tuple[str, str]]) -> List[dict]:
     history_messages = [{"role": "user", "content": m[0]} for m in history]
@@ -34,7 +38,7 @@ def create_formatted_history(history_messages: List[dict]) -> List[Tuple[str, st
             user_messages = []
             assistant_messages = []
 
-    # append any remaining messages
+    # Append any remaining messages
     if user_messages:
         formatted_history.append(("".join(user_messages), None))
     elif assistant_messages:
@@ -49,26 +53,39 @@ def chat(message: str, state: List[Dict[str, str]]) -> Generator[Tuple[List[Tupl
         history_messages.append({"role": "system", "content": "A helpful assistant."})
 
     history_messages.append({"role": "user", "content": message})
-    # We have no content for the assistant's response yet but we will update this:
     history_messages.append({"role": "assistant", "content": ""})
 
-    # Prepare input for the model
-    inputs = tokenizer(message, return_tensors="pt").to(device)
+    # Tokenize user input and prepare input tensor
+    input_ids = tokenizer.encode(message, return_tensors='pt').to(device)
+
+    if input_ids.size(-1) == 0:
+        response_message = "Input was empty after tokenization. Please try again."
+    else:
+        # Generate tokens one by one
+        with torch.no_grad():
+            for _ in range(50):  # Limit generation to 50 tokens
+                outputs = model(input_ids)
+                next_token_logits = outputs.logits[:, -1, :]
+                next_token_id = torch.argmax(next_token_logits, dim=-1)
+                input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
+
+                # Decode and append the latest token
+                decoded_token = tokenizer.decode(next_token_id)
+                history_messages[-1]["content"] += decoded_token
 
-    # Generate response from model
-    response_ids = model.generate(inputs['input_ids'], max_length=200)
-    response_message = tokenizer.decode(response_ids[0], skip_special_tokens=True)
+                # Stop if the model generates the end-of-sequence token
+                if next_token_id.item() == tokenizer.eos_token_id:
+                    break
 
-    # Update the assistant's response in our model
-    history_messages[-1]["content"] = response_message
+        response_message = history_messages[-1]["content"]
 
     formatted_history = create_formatted_history(history_messages)
     yield formatted_history, history_messages
 
-chatbot = gr.Chatbot(label="Chat").style(color_map=("yellow", "purple"))
+chatbot = gr.Chatbot(label="Chat")
 iface = gr.Interface(
     fn=chat,
-    inputs=[gr.Textbox(placeholder="Hello! How are you? etc.", label="Message"), "state"],
+    inputs=[gr.Textbox(placeholder="Hello! How are you?", label="Message"), "state"],
     outputs=[chatbot, "state"],
     allow_flagging="never",
 )