File size: 3,512 Bytes
0aeafaa 7a4f631 82d6de9 7a4f631 0aeafaa 7a4f631 82d6de9 0aeafaa 7a4f631 3b01c61 7a4f631 3b01c61 7a4f631 02a0e81 7a4f631 02a0e81 7a4f631 3b01c61 7a4f631 3b01c61 7a4f631 3b01c61 7a4f631 3b01c61 7a4f631 82d6de9 7a4f631 3b01c61 7a4f631 0aeafaa 7a4f631 82d6de9 7a4f631 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import torch
import re
# --- SETTINGS ---
MODEL_ID = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
DEVICE = "cpu"
st.set_page_config(page_title="Smol Claude", layout="wide", page_icon="🤖")
# --- CUSTOM CSS FOR CLAUDE VIBE ---
st.markdown("""
<style>
.stChatMessage { border-radius: 10px; padding: 10px; margin-bottom: 10px; }
.artifact-container { background-color: #f8f9fa; border-left: 2px solid #ddd; padding: 20px; height: 100vh; overflow-y: auto; }
</style>
""", unsafe_allow_html=True)
# --- MODEL LOADING ---
@st.cache_resource
def load_llm():
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto").to(DEVICE)
return tokenizer, model
tokenizer, model = load_llm()
# --- HELPER: EXTRACT CODE ---
def extract_code(text):
match = re.search(r"```(?:\w+)?\n([\s\S]*?)\n```", text)
return match.group(1) if match else None
# --- UI SETUP ---
if "messages" not in st.session_state:
st.session_state.messages = []
if "current_code" not in st.session_state:
st.session_state.current_code = ""
col_chat, col_artifact = st.columns([1, 1])
# --- CHAT COLUMN ---
with col_chat:
st.subheader("💬 Chat")
for msg in st.session_state.messages:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
if prompt := st.chat_input("What should we build?"):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
# Prepare conversation
input_msgs = [{"role": "system", "content": "You are a professional coder. Output code in markdown blocks."}] + st.session_state.messages
text = tokenizer.apply_chat_template(input_msgs, tokenize=False, add_generation_prompt=True)
inputs = tokenizer([text], return_tensors="pt").to(DEVICE)
# Streaming setup
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, temperature=0.1)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
# Stream the response
full_response = ""
message_placeholder = st.empty()
for new_text in streamer:
full_response += new_text
message_placeholder.markdown(full_response + "▌")
message_placeholder.markdown(full_response)
st.session_state.messages.append({"role": "assistant", "content": full_response})
# Update Artifact
code_snippet = extract_code(full_response)
if code_snippet:
st.session_state.current_code = code_snippet
st.rerun()
# --- ARTIFACT COLUMN (The "Claude" Sidebar) ---
with col_artifact:
st.subheader("🚀 Code Artifact")
if st.session_state.current_code:
st.code(st.session_state.current_code, language="python", line_numbers=True)
st.download_button("Download File", st.session_state.current_code, file_name="artifact.py")
else:
st.info("Ask the AI to write some code to see it appear here.") |