Update app.py

app.py

@@ -1,57 +1,33 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import pipeline
 
-# Load the model and tokenizer
-model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype="auto",
-    device_map="auto"
-)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Initialize the pipeline for code generation and assistance
+pipe = pipeline("text-generation", model="Qwen/Qwen2.5-Coder-32B-Instruct")
 
-# Function to interact with the model
-def chat_with_model(user_input):
-    prompt = user_input
-    messages = [
-        {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
-        {"role": "user", "content": prompt}
-    ]
-    
-    # Use apply_chat_template to format messages for the model
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    
-    # Tokenize the input and send it to the model
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+# Function to interact with the model for code-related assistance
+def code_assistance(user_input):
+    # Define the system message to set the context
+    system_message = "You are Qwen, a code assistant created by Alibaba Cloud. You assist with code generation, debugging, and explanation tasks."
     
-    # Generate the response from the model
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=512
-    )
+    # Format the prompt with the system message and user input (code-related query)
+    prompt = f"{system_message}\nUser: {user_input}\nAssistant (Code Assistance):"
     
-    # Decode the generated response
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
+    # Use the pipeline to generate the response for code assistance
+    response = pipe(prompt, max_length=512, num_return_sequences=1)
     
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    # Extract and clean the response to return only the assistant's code suggestion or explanation
+    generated_response = response[0]['generated_text'].split("Assistant (Code Assistance):")[1].strip()
     
-    return response
+    return generated_response
 
-# Create the Gradio interface
+# Create the Gradio interface for the code assistance chatbot
 iface = gr.Interface(
-    fn=chat_with_model, 
-    inputs=gr.Textbox(lines=2, placeholder="Ask me anything..."), 
+    fn=code_assistance, 
+    inputs=gr.Textbox(lines=5, placeholder="Ask for code help..."), 
     outputs="text", 
     title="Qwen2.5-Coder Chatbot", 
-    description="A chatbot using Qwen2.5-Coder for code generation, reasoning, and fixing tasks."
+    description="A chatbot using Qwen2.5-Coder for code generation, debugging, and explanation tasks."
 )
 
-# Launch the interface
-iface.launch()
+# Launch the Gradio interface
+iface.launch()