Update app.py

app.py

@@ -1,31 +1,59 @@
+import os
+os.system("pip install transformers")
+
+import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
-device = "cuda" # the device to load the model onto
 
+# Set device
+device = "cuda" if st.sidebar.checkbox("Use GPU", True) else "cpu"
+
+# Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
     "Qwen/Qwen1.5-0.5B-Chat",
     torch_dtype="auto",
     device_map="auto"
-)
+).to(device)
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
 
-prompt = "Give me a short introduction to large language model."
+# Create a chatbot interface
+st.title("Chatbot")
+st.write("Ask me anything!")
+
+# Initialize messages
 messages = [
     {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": prompt}
-]
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(device)
-
-generated_ids = model.generate(
-    model_inputs.input_ids,
-    max_new_tokens=512
-)
-generated_ids = [
-    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
 ]
 
-response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+# Display chat history
+for message in messages:
+    if message["role"] == "system":
+        st.write(f"*System*: {message['content']}")
+    elif message["role"] == "user":
+        st.write(f"*You*: {message['content']}")
+    elif message["role"] == "assistant":
+        st.write(f"*Assistant*: {message['content']}")
+
+# Get user input
+user_input = st.text_input("Your message")
+
+# Generate response
+if user_input:
+    messages.append({"role": "user", "content": user_input})
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    model_inputs = tokenizer([text], return_tensors="pt").to(device)
+    generated_ids = model.generate(
+        model_inputs.input_ids,
+        max_new_tokens=512
+    )
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    messages.append({"role": "assistant", "content": response})
+
+    # Display response
+    st.write(f"*Assistant*: {response}")