File size: 4,439 Bytes
5b2a0f1 dc7ac95 5b2a0f1 dc7ac95 5b2a0f1 dc7ac95 5b2a0f1 dc7ac95 5b2a0f1 dc7ac95 5b2a0f1 dc7ac95 5b2a0f1 |
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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import streamlit as st
import torch
from model_handler import load_model_and_tokenizer, generate_response
import time
# Page configuration
st.set_page_config(
page_title="Phi-2 Fine-tuned Assistant",
page_icon="🤖",
layout="centered"
)
# Custom CSS for chat interface
st.markdown("""
<style>
.user-bubble {
background-color: #e6f7ff;
border-radius: 15px;
padding: 10px 15px;
margin: 5px 0;
max-width: 80%;
margin-left: auto;
margin-right: 10px;
position: relative;
color: #000000; /* Ensure text is black */
}
.assistant-bubble {
background-color: #f0f0f0;
border-radius: 15px;
padding: 10px 15px;
margin: 5px 0;
max-width: 80%;
margin-left: 10px;
position: relative;
color: #000000; /* Ensure text is black */
}
.chat-container {
display: flex;
flex-direction: column;
height: calc(100vh - 200px);
overflow-y: auto;
padding: 10px;
margin-bottom: 20px;
}
.stTextInput>div>div>input {
border-radius: 20px;
}
.stButton>button {
border-radius: 20px;
width: 100%;
}
/* Fix for excessive vertical space */
header {
visibility: hidden;
}
.block-container {
padding-top: 1rem;
padding-bottom: 1rem;
}
h1 {
margin-top: 0 !important;
margin-bottom: 1rem !important;
}
/* Ensure dark mode compatibility */
.stApp {
background-color: #121212;
}
</style>
""", unsafe_allow_html=True)
# Initialize session state for chat history
if "messages" not in st.session_state:
st.session_state.messages = []
if "model_loaded" not in st.session_state:
st.session_state.model_loaded = False
# App title - made smaller to reduce vertical space
st.markdown("<h2>Phi-2 Fine-tuned Assistant</h2>", unsafe_allow_html=True)
# Load model (only once)
if not st.session_state.model_loaded:
with st.spinner("Loading the fine-tuned model... This may take a minute."):
model, tokenizer = load_model_and_tokenizer()
st.session_state.model = model
st.session_state.tokenizer = tokenizer
st.session_state.model_loaded = True
st.success("Model loaded successfully!")
# Display chat messages
st.markdown('<div class="chat-container">', unsafe_allow_html=True)
for message in st.session_state.messages:
role = message["role"]
content = message["content"]
if role == "user":
st.markdown(f'<div class="user-bubble">{content}</div>', unsafe_allow_html=True)
else:
st.markdown(f'<div class="assistant-bubble">{content}</div>', unsafe_allow_html=True)
st.markdown('</div>', unsafe_allow_html=True)
# Chat input
with st.container():
col1, col2 = st.columns([5, 1])
with col1:
user_input = st.text_input("", key="user_input", placeholder="Type your message here...")
with col2:
clear_button = st.button("Clear")
# Process input when user submits a message
if user_input:
# Add user message to chat history
st.session_state.messages.append({"role": "user", "content": user_input})
# Create a placeholder for the assistant's response
assistant_response_placeholder = st.empty()
# Generate streaming response
full_response = ""
# Display "Assistant is typing..." message
with assistant_response_placeholder.container():
st.markdown('<div class="assistant-bubble">Assistant is typing...</div>', unsafe_allow_html=True)
# Generate response with streaming
for token in generate_response(
st.session_state.model,
st.session_state.tokenizer,
user_input # Changed to pass just the current input
):
full_response += token
# Update the response in real-time
with assistant_response_placeholder.container():
st.markdown(f'<div class="assistant-bubble">{full_response}</div>', unsafe_allow_html=True)
time.sleep(0.01) # Small delay to make streaming visible
# Add assistant's response to chat history
st.session_state.messages.append({"role": "assistant", "content": full_response})
# Clear the input box
st.session_state.user_input = ""
# Rerun to update the UI
st.experimental_rerun()
# Clear chat when button is pressed
if clear_button:
st.session_state.messages = []
st.experimental_rerun()
|