iforgebot / app.py
tenet's picture
Update app.py
c3ab621 verified
import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode
# Load the pretrained DialoGPT model
tokenizer = GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium")
# Streamlit UI Setup
st.title("AI Multimodal Chat & File Processing App")
# Session state for chat history
if "history" not in st.session_state:
st.session_state.history = []
# Function to process the chat
def chat_with_model(user_input):
new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
st.session_state.history.append(new_user_input_ids)
bot_input_ids = new_user_input_ids
for history in st.session_state.history:
bot_input_ids = history if len(history) < 2048 else history[-1024:]
chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
bot_output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
return bot_output
# Chat Input Box
st.subheader("Chat with AI")
user_input = st.text_input("You: ", "")
if user_input:
response = chat_with_model(user_input)
st.session_state.history.append(tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt"))
st.write(f"Bot: {response}")
# Show chat history (safer decoding)
if st.session_state.history:
st.subheader("πŸ“ Chat History")
for i in range(len(st.session_state.history) - 1, -1, -1):
try:
user_msg = tokenizer.decode(st.session_state.history[i][0].tolist(), skip_special_tokens=True)
st.write(f"You: {user_msg}")
except Exception as e:
st.warning(f"Could not decode history message: {e}")
# --- File Upload and Processing ---
st.subheader("πŸ“ Upload a File for AI to Read")
uploaded_file = st.file_uploader("Choose a text file", type=["txt", "csv", "md", "log"])
if uploaded_file:
content = uploaded_file.read().decode("utf-8")
st.text_area("File Content", content, height=200)
# Allow interaction with file content
file_question = st.text_input("Ask something about the file:")
if file_question:
combined_input = file_question + "\n" + content[:1000] # Prevent token overload
response = chat_with_model(combined_input)
st.write(f"Bot: {response}")
# --- Video/Audio Stream ---
st.subheader("πŸŽ₯ Video/Audio Stream")
class VideoProcessor(VideoProcessorBase):
def recv(self, frame):
return frame
webrtc_streamer(key="example", mode=WebRtcMode.SENDRECV, video_processor_factory=VideoProcessor)