Update app.py

app.py

@@ -1,38 +1,40 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the AIDC-AI/Marco-o1 model and tokenizer
+# Load the Marco-o1 model and tokenizer
 model_name = "AIDC-AI/Marco-o1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Function to handle conversation between two AI models
+# Conversation function
 def ai_conversation(prompt, history=[]):
-    # AI 1 responds
+    # Ensure the history starts properly
+    if history is None:
+        history = []
+    
+    # AI 1 generates a response
     input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
     response_ids = model.generate(input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id)
     response1 = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    
     history.append(("You", prompt))
     history.append(("AI 1", response1))
     
-    # AI 2 responds
+    # AI 2 responds to AI 1
     input_ids = tokenizer.encode(response1 + tokenizer.eos_token, return_tensors="pt")
     response_ids = model.generate(input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id)
     response2 = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    
     history.append(("AI 1", response1))
     history.append(("AI 2", response2))
 
     return history, history
 
-# Create Gradio Interface
+# Gradio Interface
 interface = gr.Interface(
     fn=ai_conversation,
     inputs=["text", "state"],
     outputs=["chatbot", "state"],
-    title="Marco-o1 AI-to-AI Conversation",
-    description="Type a prompt to initiate a conversation between two instances of Marco-o1."
+    title="Marco-o1 Group Chat Simulation",
+    description="Type a message to start a group chat between two AI instances."
 )
 
 if __name__ == "__main__":