Update app.py

app.py

@@ -1,6 +1,6 @@
 import re
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, TFAutoModelForMaskedLM
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
 # Initialize the chat history
 history = []
@@ -8,8 +8,8 @@ history = []
 def clean_text(text):
     return re.sub('[^a-zA-Z\s]', '', text).strip()
 
-tokenizer = AutoTokenizer.from_pretrained("t5-small")
-model = TFAutoModelForMaskedLM.from_pretrained("t5-small").half().cuda()
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
+model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/DialoGPT-small").half().cuda()
 
 def generate_response(user_input):
     history.append((user_input, ""))
@@ -18,27 +18,28 @@ def generate_response(user_input):
         return ""
     
     last_user_message = history[-1][0]
-    combined_messages = " Human: " + " . ".join([msg for msg, _ in reversed(history[:-1])]) + " . Human: " + last_user_message
-    input_str = "summarize: " + combined_messages
-    source_encodings = tokenizer.batch_encode_plus([input_str], pad_to_max_length=False, padding='max_length', return_attention_mask=True, return_tensors="tf")
-    input_ids = source_encodings["input_ids"][0]
-    attention_mask = source_encodings["attention_mask"][0]
-    input_ids = tf.constant(input_ids)[None, :]
-    attention_mask = tf.constant(attention_mask)[None, :]
-    
-    with tf.device('/GPU:0'):
-      output = model.generate(
-          input_ids,
-          attention_mask=attention_mask,
-          max_length=256,
-          num_beams=4,
-          early_stopping=True
-      )
-      
-    predicted_sentence = tokenizer.decode(output[0], skip_special_tokens=True)
+    combined_messages = " ".join([msg for msg, _ in reversed(history[:-1])]) + " User: " + last_user_message
+
+    tokens = tokenizer.encode(combined_messages, add_special_tokens=True, max_length=4096, truncation=True)
+    tokens = tokens[:1024]
+    segment_ids = [0]*len(tokens)
+    input_ids = torch.tensor([tokens], dtype=torch.long).cuda()
+
+    with torch.no_grad():
+        outputs = model.generate(
+           input_ids,
+           max_length=1024,
+           min_length=20,
+           length_penalty=2.0,
+           early_stopping=True,
+           num_beams=4,
+           bad_words_callback=[lambda x: True if 'User:' in str(x) else False]
+        )
+    output = output[0].tolist()[len(tokens)-1:]
+    decoded_output = tokenizer.decode(output, skip_special_tokens=True)
 
-    history[-1] = (last_user_message, predicted_sentence)
-    return f"AI: {predicted_sentence}".capitalize()
+    history[-1] = (last_user_message, decoded_output)
+    return f"AI: {decoded_output}".capitalize()
 
 st.title("Simple Chat App using DistilBert Model (HuggingFace & Streamlit)")