Update app.py

app.py

@@ -1,69 +1,104 @@
 # app.py
 # =============
-# This is a complete app.py file for a Gradio app using the meta-llama/Llama-3.2-3B-Instruct model.
-# The app allows users to input a message and receive a response from the model.
+# This is a complete app.py file for a text generation app using the Qwen/Qwen2.5-Coder-0.5B-Instruct model.
+# The app uses the Gradio library to create a web interface for interacting with the model.
 
-# Dependencies
-# =============
-# The following dependencies are required to run this app:
-# - transformers
-# - gradio
-# - torch
-#
-# You can install these dependencies using pip:
-# pip install transformers gradio torch
+# Imports
+# =======
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-import torch
-from transformers import pipeline
-import gradio as gr  # Import gradio
+# Constants
+# =========
+MODEL_NAME = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+SYSTEM_MESSAGE = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
 
-# Load the model and tokenizer
-model_id = "meta-llama/Llama-3.2-3B-Instruct"
-device = "cpu"  # Use CPU for inference
+# Load Model and Tokenizer
+# ========================
+def load_model_and_tokenizer():
+    """
+    Load the model and tokenizer from Hugging Face.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype="auto",
+        device_map="cpu"  # Ensure the model runs on the CPU
+    )
+    return model, tokenizer
 
-# Initialize the pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
+model, tokenizer = load_model_and_tokenizer()
 
-def generate_response(prompt):
+# Generate Response
+# =================
+def generate_response(prompt, chat_history):
     """
-    Generate a response from the model based on the given prompt.
-    
-    Args:
-        prompt (str): The input message from the user.
-    
-    Returns:
-        str: The generated response from the model.
+    Generate a response from the model based on the user prompt and chat history.
     """
-    messages = [
-        {"role": "system", "content": "You are a helpful assistant!"},
-        {"role": "user", "content": prompt},
-    ]
-    outputs = pipe(
+    messages = [{"role": "system", "content": SYSTEM_MESSAGE}] + chat_history + [{"role": "user", "content": prompt}]
+    text = tokenizer.apply_chat_template(
         messages,
-        max_new_tokens=256,
+        tokenize=False,
+        add_generation_prompt=True
     )
-    return outputs[0]["generated_text"][-1]
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=512,
+        do_sample=True,
+        top_k=50,
+        top_p=0.95,
+        temperature=0.7,
+        stream=True
+    )
+
+    response = ""
+    for new_token in generated_ids[0][len(model_inputs.input_ids[0]):]:
+        response += tokenizer.decode([new_token], skip_special_tokens=True)
+        yield response
+
+# Clear Chat History
+# ==================
+def clear_chat():
+    """
+    Clear the chat history.
+    """
+    return [], []
 
-# Define the Gradio interface
+# Gradio Interface
+# =================
 def gradio_interface():
     """
-    Define the Gradio interface for the app.
+    Create and launch the Gradio interface.
     """
-    iface = gr.Interface(
-        fn=generate_response,
-        inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your message here..."),
-        outputs="text",
-        title="Llama-3.2-3B-Instruct Chatbot",
-        description="Chat with the Llama-3.2-3B-Instruct model. Enter your message and get a response!",
-    )
-    return iface
+    with gr.Blocks() as demo:
+        chatbot = gr.Chatbot(label="Chat with Qwen/Qwen2.5-Coder-0.5B-Instruct")
+        msg = gr.Textbox(label="User Input")
+        clear = gr.Button("Clear Chat")
+
+        def respond(message, chat_history):
+            chat_history.append({"role": "user", "content": message})
+            response = generate_response(message, chat_history)
+            chat_history.append({"role": "assistant", "content": response})
+            return chat_history, chat_history
 
-# Launch the Gradio app
+        msg.submit(respond, [msg, chatbot], [chatbot, chatbot])
+        clear.click(clear_chat, None, [chatbot, chatbot])
+
+    demo.launch()
+
+# Main
+# ====
 if __name__ == "__main__":
-    iface = gradio_interface()
-    iface.launch()
+    gradio_interface()
+
+# Dependencies
+# =============
+# The following dependencies are required to run this app:
+# - transformers
+# - gradio
+# - torch
+#
+# You can install these dependencies using pip:
+# pip install transformers gradio torch