iforgebot / app.py
tenet's picture
Update app.py
24c37ee verified
raw
history blame
1.88 kB
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")
# Chat history session state setup
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
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
if st.session_state.history:
for i in range(len(st.session_state.history) - 1, -1, -1):
user_msg = tokenizer.decode(st.session_state.history[i], skip_special_tokens=True)
st.write(f"You: {user_msg}")
# 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)