mew77's picture
Update app.py
a71142d verified
import streamlit as st
from PIL import Image
import os
from datetime import datetime
from hf_model import HFModel # Import your HFModel class
from mew_log import log_info, log_error # Import your custom logging methods
# Initialize the model
MODEL_NAME = "deepseek-ai/DeepSeek-Coder-V2-Lite-Base"
def initialize_model():
try:
model = HFModel(MODEL_NAME)
log_info("Model initialized successfully.")
return model
except Exception as e:
log_error(f"Failed to initialize model: {e}", e)
st.error(f"Failed to initialize model: {e}")
return None
# Initialize the model (cached to avoid reloading on every interaction)
@st.cache_resource
def load_model():
return initialize_model()
# Streamlit app
def main():
st.title("DeepSeek Chat App")
st.write("Interact with the DeepSeek model using text, images, or files!")
# Load the model
model = load_model()
if model is None:
st.stop()
# Sidebar for file uploads
st.sidebar.header("Upload Files or Images")
uploaded_file = st.sidebar.file_uploader("Upload a file or image", type=["txt", "jpg", "png", "jpeg", "pdf"])
# Display uploaded file or image
if uploaded_file is not None:
file_details = {"filename": uploaded_file.name, "filetype": uploaded_file.type, "filesize": uploaded_file.size}
st.sidebar.write(file_details)
if uploaded_file.type.startswith("image"):
image = Image.open(uploaded_file)
st.sidebar.image(image, caption="Uploaded Image", use_column_width=True)
else:
st.sidebar.write("File content:")
file_content = uploaded_file.getvalue().decode("utf-8")
st.sidebar.text(file_content)
# Main chat interface
st.header("Chat with DeepSeek")
user_input = st.text_area("Enter your message or question:", height=100)
if st.button("Send"):
if user_input.strip() == "":
st.warning("Please enter a message.")
else:
try:
# Add file content to user input if a file is uploaded
if uploaded_file is not None:
if uploaded_file.type.startswith("image"):
user_input += f"\n[Attached Image: {uploaded_file.name}]"
else:
file_content = uploaded_file.getvalue().decode("utf-8")
user_input += f"\n[Attached File Content:\n{file_content}\n]"
# Generate response
with st.spinner("Generating response..."):
response = model.chat(user_input, max_length=512)
st.write(f"**Assistant:** {response}")
# Display chat history
st.subheader("Chat History")
for entry in model.chat_history:
role = entry["role"]
content = entry["content"]
st.write(f"**{role.capitalize()}:** {content}")
except Exception as e:
log_error(f"Error during chat: {e}", e)
st.error(f"An error occurred: {e}")
# Clear chat history button
if st.button("Clear Chat History"):
model.clear_chat_history()
st.success("Chat history cleared!")
# Run the app
if __name__ == "__main__":
main()