Update app.py

app.py

@@ -3,31 +3,23 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 ### Model Setup ###
-# Names for the two models
-LLAMA_MODEL_NAME = "meta-llama/Llama-3.2-11B-Vision"
-QWEN_MODEL_NAME = "Qwen/Qwen2.5-VL-7B-Instruct"
+# Define the two model names.
+MODEL1_NAME = "Qwen/Qwen2.5-14B-Instruct-1M"   # We'll refer to this as Qwen14B.
+MODEL2_NAME = "Qwen/Qwen2.5-VL-7B-Instruct"       # We'll refer to this as QwenVL.
 
-# Load Meta-Llama Vision model
-llama_tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME)
-llama_model = AutoModelForCausalLM.from_pretrained(LLAMA_MODEL_NAME, device_map="auto")
-llama_pipe = pipeline(
-    "text-generation", 
-    model=llama_model, 
-    tokenizer=llama_tokenizer
-)
+# Load Qwen14B model.
+tokenizer1 = AutoTokenizer.from_pretrained(MODEL1_NAME)
+model1 = AutoModelForCausalLM.from_pretrained(MODEL1_NAME, device_map="auto")
+pipe1 = pipeline("text-generation", model=model1, tokenizer=tokenizer1)
 
-# Load Qwen
-qwen_tokenizer = AutoTokenizer.from_pretrained(QWEN_MODEL_NAME)
-qwen_model = AutoModelForCausalLM.from_pretrained(QWEN_MODEL_NAME, device_map="auto")
-qwen_pipe = pipeline(
-    "text-generation", 
-    model=qwen_model, 
-    tokenizer=qwen_tokenizer
-)
+# Load QwenVL model.
+tokenizer2 = AutoTokenizer.from_pretrained(MODEL2_NAME)
+model2 = AutoModelForCausalLM.from_pretrained(MODEL2_NAME, device_map="auto")
+pipe2 = pipeline("text-generation", model=model2, tokenizer=tokenizer2)
 
 def generate_response(prompt: str, model_choice: str) -> str:
     """
-    Given a prompt and a model choice, generate a response.
+    Generate a response given the conversation prompt using the chosen model.
     """
     kwargs = {
         "max_length": 256,
@@ -35,40 +27,38 @@ def generate_response(prompt: str, model_choice: str) -> str:
         "top_p": 0.95,
         "temperature": 0.9,
     }
-    if model_choice == "Llama":
-        result = llama_pipe(prompt, **kwargs)
-    else:  # "Qwen"
-        result = qwen_pipe(prompt, **kwargs)
-    # Extract and return the generated text
+    if model_choice == "Qwen14B":
+        result = pipe1(prompt, **kwargs)
+    else:  # model_choice == "QwenVL"
+        result = pipe2(prompt, **kwargs)
     return result[0]["generated_text"]
 
 def chat_logic(user_input: str, chat_history: list):
     """
-    Build the conversation prompt from the history, decide which model responds,
-    and update the conversation.
+    Build the conversation prompt from the history, choose which model responds,
+    generate the response, and update the conversation.
     """
     if chat_history is None:
         chat_history = []
     
-    # If the user provided input, record it and choose a random model to respond.
+    # If the user provides input, add it to the history and randomly choose a model.
     if user_input.strip():
         chat_history.append(("User", user_input.strip()))
-        selected_model = random.choice(["Llama", "Qwen"])
+        selected_model = random.choice(["Qwen14B", "QwenVL"])
     else:
-        # If no user input, let the models take turns.
+        # When no user input is provided, let the models alternate.
         if not chat_history:
-            selected_model = random.choice(["Llama", "Qwen"])
+            selected_model = random.choice(["Qwen14B", "QwenVL"])
         else:
-            # If the last speaker was a model, switch to the other model.
             last_speaker = chat_history[-1][0]
-            if last_speaker == "Llama":
-                selected_model = "Qwen"
-            elif last_speaker == "Qwen":
-                selected_model = "Llama"
+            if last_speaker == "Qwen14B":
+                selected_model = "QwenVL"
+            elif last_speaker == "QwenVL":
+                selected_model = "Qwen14B"
             else:
-                selected_model = random.choice(["Llama", "Qwen"])
+                selected_model = random.choice(["Qwen14B", "QwenVL"])
     
-    # Construct the conversation prompt.
+    # Build the prompt from the conversation history.
     prompt = ""
     for speaker, message in chat_history:
         prompt += f"{speaker}: {message}\n"
@@ -78,38 +68,38 @@ def chat_logic(user_input: str, chat_history: list):
     model_response = generate_response(prompt, selected_model)
     chat_history.append((selected_model, model_response))
     
-    # Return both the updated chat for display and as state.
+    # Return the updated conversation for display and to maintain state.
     return chat_history, chat_history
 
 ### Gradio Interface ###
 with gr.Blocks() as demo:
-    gr.Markdown("# Group Chat: Meta-Llama Vision and Qwen Models")
+    gr.Markdown("# Group Chat: Qwen Models")
     gr.Markdown(
-        "This demo allows two models to converse with each other. "
-        "Leave the textbox blank to let the models continue the conversation automatically, "
-        "or type a message to interject (a random model will then respond)."
+        "This demo features two Qwen models conversing with each other. "
+        "Leave the textbox blank to let the models converse automatically, "
+        "or type a message to interject (a random model will then reply)."
     )
     
-    # Chatbot display component.
+    # Chat display component.
     chatbot = gr.Chatbot()
     
-    # Input row: text input and a send button.
+    # Input row: textbox and send button.
     with gr.Row():
         user_message = gr.Textbox(
-            placeholder="Type a message or leave blank to continue conversation...",
+            placeholder="Type your message here or leave blank...",
             show_label=False
         )
         send_btn = gr.Button("Send")
     
-    # Maintain the conversation history in Gradio's state.
+    # Maintain the conversation history.
     state = gr.State([])
 
-    # Update the chat when the send button is clicked.
+    # On clicking the button, update the conversation.
     send_btn.click(
         fn=chat_logic,
         inputs=[user_message, state],
         outputs=[chatbot, state]
     )
 
-# Launch the Space
+# Launch the Space.
 demo.launch()