AIchat-2 / app.py
Ivan000's picture
Update app.py
baecb1b verified
raw
history blame
3.01 kB
# app.py
# =============
# 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.
# Imports
# =======
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
# 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 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
model, tokenizer = load_model_and_tokenizer()
# Generate Response
# =================
def generate_response(prompt, chat_history):
"""
Generate a response from the model based on the user prompt and chat history.
"""
messages = [{"role": "system", "content": SYSTEM_MESSAGE}] + chat_history + [{"role": "user", "content": prompt}]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
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 [], []
# Gradio Interface
# =================
def gradio_interface():
"""
Create and launch the Gradio interface.
"""
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
msg.submit(respond, [msg, chatbot], [chatbot, chatbot])
clear.click(clear_chat, None, [chatbot, chatbot])
demo.launch()
# Main
# ====
if __name__ == "__main__":
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