Update app.py

app.py

@@ -87,25 +87,65 @@
 # text = tokenizer.batch_decode(outputs)[0]
 # print(text)
 
-import torch
-import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+# import torch
+# import streamlit as st
+# from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-model_name="facebook/blenderbot-400M-distill"
+# model_name="facebook/blenderbot-400M-distill"
 
-model=AutoModelForSeq2SeqLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-ch=[]
-def chat():
+# model=AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# ch=[]
+# def chat():
     
-    h_s="\n".join(ch)
-    i=st.text_input("enter")
-    i_s=tokenizer.encode_plus(h_s,i,return_tensors="pt")
-    outputs=model.generate(**i_s,max_length=60)
-    response=tokenizer.decode(outputs[0],skip_special_tokens=True).strip()
-    ch.append(i)
-    ch.append(response)
-    return response
-if __name__ == "__main__":
-    chat()
+#     h_s="\n".join(ch)
+#     i=st.text_input("enter")
+#     i_s=tokenizer.encode_plus(h_s,i,return_tensors="pt")
+#     outputs=model.generate(**i_s,max_length=60)
+#     response=tokenizer.decode(outputs[0],skip_special_tokens=True).strip()
+#     ch.append(i)
+#     ch.append(response)
+#     return response
+# if __name__ == "__main__":
+#     chat()
+
+import streamlit as st
+from streamlit_chat import message as st_message
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+
+@st.experimental_singleton
+def get_models():
+    # Load the model and the tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot_small-90M")
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        "facebook/blenderbot_small-90M")
+
+    return tokenizer, model
+
+
+if "history" not in st.session_state:
+    st.session_state.history = []
+
+st.title("Blenderbot")
+
+
+def generate_answer():
+    tokenizer, model = get_models()
+    user_message = st.session_state.input_text
+    inputs = tokenizer(st.session_state.input_text, return_tensors="pt")
+    result = model.generate(**inputs)
+    message_bot = tokenizer.decode(
+        result[0], skip_special_tokens=True
+    )   # decode the result to a string
+
+    st.session_state.history.append({"message": user_message, "is_user": True})
+    st.session_state.history.append({"message": message_bot, "is_user": False})
+
+
+st.text_input("Tap to chat with the bot",
+              key="input_text", on_change=generate_answer)
 
+for chat in st.session_state.history:
+    st_message(**chat)