add app

app.py

@@ -1,57 +1,10 @@
-# import streamlit as st
-# from streamlit_chat import message as st_message
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 MAX_HISTORY = 7
 MODEL_PATH = 'llongpre/DialoGPT-small-miles'
 
-def get_models():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
-    return tokenizer, model
-
-# if "history" not in st.session_state:
-#     st.session_state.history = []
-#
-# if "history_ids" not in st.session_state:
-#     st.session_state.history_ids = []
-#
-# st.title("Chat with me")
-
-# def generate_answer():
-#     tokenizer, model = get_models()
-#     user_message = st.session_state.input_text
-#     new_user_input_ids = tokenizer.encode(st.session_state.input_text + tokenizer.eos_token, return_tensors='pt')
-#     st.session_state.history_ids.append(new_user_input_ids)
-#     if len(st.session_state.history_ids) > MAX_HISTORY:
-#         st.session_state.history_ids = st.session_state.history_ids[-MAX_HISTORY:]
-#     bot_input_ids = torch.cat(st.session_state.history_ids, dim=-1)
-#     chat_history_ids = model.generate(
-#         bot_input_ids,
-#         pad_token_id=tokenizer.pad_token_id,
-#         max_length=1000,
-#         do_sample=True,
-# #         top_k=150, # sample from the top k words sorted descending by probability
-#         top_p=0.7, # choose smallest possible words whose cumulative probability exceeds p
-#         temperature = 0.95, # 0 greedy, inf is random
-#         no_repeat_ngram_size=3,
-#     )
-#     response = chat_history_ids[:, bot_input_ids.shape[-1]:]
-#     st.session_state.history_ids.append(response)
-#     output = tokenizer.decode(response[0], skip_special_tokens=True)
-#
-#     st.session_state.history.append({"message": user_message, "is_user": True})
-#     st.session_state.history.append({"message": output, "is_user": False})
-
-# st.text_input("Your text message", key="input_text", on_change=generate_answer, placeholder='')
-
-# for chat in st.session_state.history:
-#     st_message(**chat)  # unpacking
-
-
-import gradio as gr
-
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
 
@@ -80,8 +33,7 @@ def predict(input, history=[]):
 
 def generate_answer(input, history=[]):
     new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors='pt')
-    # st.session_state.history_ids.append(new_user_input_ids)
-    history = history.append(new_user_input_ids)
+    history = history.append(input)
     if len(history) > MAX_HISTORY:
         history = history[-MAX_HISTORY:]
     bot_input_ids = torch.cat(history, dim=-1)
@@ -96,8 +48,8 @@ def generate_answer(input, history=[]):
         no_repeat_ngram_size=3,
     )
     response = chat_history_ids[:, bot_input_ids.shape[-1]:]
-    history.append(response)
     output = tokenizer.decode(response[0], skip_special_tokens=True)
+    history.append(output)
 
     return output, history