Spaces:
Sleeping
Sleeping
File size: 3,120 Bytes
967d7c5 8c113c3 967d7c5 e74f40b 967d7c5 f3c6eb8 967d7c5 6d9b383 967d7c5 9dc275b 967d7c5 9dc275b 967d7c5 8c113c3 967d7c5 9dc275b 967d7c5 f3c6eb8 967d7c5 e74f40b 6d9b383 ddf509e 9dc275b 812346b 6d9b383 8c113c3 7087b82 9dc275b 6d9b383 e74f40b f3c6eb8 e74f40b 6d9b383 9dc275b 967d7c5 9dc275b 967d7c5 6d9b383 f3c6eb8 e74f40b 967d7c5 8c113c3 e74f40b 967d7c5 6d9b383 967d7c5 e74f40b 967d7c5 9dc275b 6d9b383 9dc275b 967d7c5 9dc275b 967d7c5 6d9b383 e74f40b 9dc275b 967d7c5 6d9b383 9dc275b 967d7c5 9dc275b e74f40b 8c113c3 e74f40b 8c113c3 e74f40b 8c113c3 e74f40b | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import streamlit as st
import torch
import re
from transformers import AutoTokenizer, AutoModelForCausalLM
# ==============================
# PAGE CONFIG
# ==============================
st.set_page_config(page_title="🤖 AI Assistant", layout="wide")
st.title("🤖 Simple AI Assistant")
# ==============================
# LOAD MODEL (CHAT MODEL ✅)
# ==============================
@st.cache_resource
def load_model():
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # ✅ Best for HF free
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float32
)
model.eval()
return tokenizer, model
with st.spinner("🔄 Loading model..."):
tokenizer, model = load_model()
st.success("✅ Ready")
# ==============================
# SESSION STATE
# ==============================
if "messages" not in st.session_state:
st.session_state.messages = []
# ==============================
# CLEAN TEXT
# ==============================
def clean_text(text):
text = re.sub(r"[^\x00-\x7F]+", "", text).strip()
# Ensure response completes nicely
if not text.endswith((".", "!", "?")):
text += "..."
return text
# ==============================
# GENERATE RESPONSE
# ==============================
def generate_response(user_input):
prompt = f"""
<|user|>
{user_input}
Give a clear and complete answer.
<|assistant|>
"""
inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=300, # 🔥 prevents cut-off
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.eos_token_id
)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract assistant response
if "<|assistant|>" in result:
result = result.split("<|assistant|>")[-1]
return clean_text(result)
# ==============================
# DISPLAY CHAT
# ==============================
for msg in st.session_state.messages:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
# ==============================
# INPUT BOX
# ==============================
user_input = st.chat_input("Type your message...")
if user_input:
# Add user message
st.session_state.messages.append({"role": "user", "content": user_input})
with st.chat_message("user"):
st.markdown(user_input)
# Generate response
with st.spinner("🤖 Thinking..."):
response = generate_response(user_input)
# Add assistant response
st.session_state.messages.append({"role": "assistant", "content": response})
with st.chat_message("assistant"):
st.markdown(response)
# ==============================
# SIDEBAR
# ==============================
st.sidebar.title("⚙️ Options")
if st.sidebar.button("🧹 Clear Chat"):
st.session_state.messages = [] |