Update app.py

app.py

@@ -1,21 +1,36 @@
 import streamlit as st
-from transformers import pipeline
-pipe = pipeline('text-generation', model='Pclanglais/MonadGPT', device='cuda')
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
+# Load DialoGPT model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 
-def main():
-  st.title("MonadGPT Streamlit App")
-  st.markdown("This is a streamlit app that allows you to have a conversation abbout any topic")
+# Function to generate a response
+def generate_response(chat_history_ids, new_user_input):
+    new_user_input_ids = tokenizer.encode(new_user_input + tokenizer.eos_token, return_tensors='pt')
+    bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids
+    return model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
 
-  user_input = st.text_input("Your question:")
+# Streamlit app
+st.title("DialoGPT Chat")
 
-  if user_input:
-    response = pipe(user_input, max_length=256, do_sample=True, top_k=50, top_p=0.95, early_stopping=True)
-    st.write(response[0]['generated_text'])
+# Initialize chat history
+chat_history_ids = None
 
-  if st.button("Restart"):
-    st.empty()
+# Conversation loop
+while st.button("Restart Conversation"):
+    # Get user input
+    user_input = st.text_input("You:", "")
+
+    if user_input:
+        # Generate response
+        chat_history_ids = generate_response(chat_history_ids, user_input)
+
+        # Display DialoGPT response
+        st.text("DialoGPT: {}".format(tokenizer.decode(chat_history_ids[:, -1][0], skip_special_tokens=True)))
+
+# Inform the user that the conversation has ended
+st.text("Conversation Ended. Press the 'Restart Conversation' button to start a new conversation.")
 
-if __name__ == "__main__":
-  main()