Update app.py
Browse files
app.py
CHANGED
@@ -1,6 +1,7 @@
|
|
1 |
import re
|
2 |
import streamlit as st
|
3 |
from transformers import pipeline
|
|
|
4 |
|
5 |
# Initialize the chat history
|
6 |
history = []
|
@@ -8,31 +9,37 @@ history = []
|
|
8 |
def clean_text(text):
|
9 |
return re.sub('[^a-zA-Z\s]', '', text).strip()
|
10 |
|
11 |
-
|
12 |
-
model =
|
13 |
|
14 |
def generate_response(user_input):
|
15 |
-
# Add user input to history
|
16 |
history.append((user_input, ""))
|
17 |
|
18 |
if not history:
|
19 |
return ""
|
20 |
|
21 |
last_user_message = history[-1][0]
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
|
35 |
-
|
|
|
36 |
|
37 |
st.title("Simple Chat App using DistilBert Model (HuggingFace & Streamlit)")
|
38 |
|
|
|
1 |
import re
|
2 |
import streamlit as st
|
3 |
from transformers import pipeline
|
4 |
+
from transformers import AutoTokenizer, TFAutoModelForMaskedLM
|
5 |
|
6 |
# Initialize the chat history
|
7 |
history = []
|
|
|
9 |
def clean_text(text):
|
10 |
return re.sub('[^a-zA-Z\s]', '', text).strip()
|
11 |
|
12 |
+
tokenizer = AutoTokenizer.from_pretrained("t5-small")
|
13 |
+
model = TFAutoModelForMaskedLM.from_pretrained("t5-small").half().cuda()
|
14 |
|
15 |
def generate_response(user_input):
|
|
|
16 |
history.append((user_input, ""))
|
17 |
|
18 |
if not history:
|
19 |
return ""
|
20 |
|
21 |
last_user_message = history[-1][0]
|
22 |
+
combined_messages = " Human: " + " . ".join([msg for msg, _ in reversed(history[:-1])]) + " . Human: " + last_user_message
|
23 |
+
input_str = "summarize: " + combined_messages
|
24 |
+
source_encodings = tokenizer.batch_encode_plus([input_str], pad_to_max_length=False, padding='max_length', return_attention_mask=True, return_tensors="tf")
|
25 |
+
input_ids = source_encodings["input_ids"][0]
|
26 |
+
attention_mask = source_encodings["attention_mask"][0]
|
27 |
+
input_ids = tf.constant(input_ids)[None, :]
|
28 |
+
attention_mask = tf.constant(attention_mask)[None, :]
|
29 |
+
|
30 |
+
with tf.device('/GPU:0'):
|
31 |
+
output = model.generate(
|
32 |
+
input_ids,
|
33 |
+
attention_mask=attention_mask,
|
34 |
+
max_length=256,
|
35 |
+
num_beams=4,
|
36 |
+
early_stopping=True
|
37 |
+
)
|
38 |
+
|
39 |
+
predicted_sentence = tokenizer.decode(output[0], skip_special_tokens=True)
|
40 |
|
41 |
+
history[-1] = (last_user_message, predicted_sentence)
|
42 |
+
return f"AI: {predicted_sentence}".capitalize()
|
43 |
|
44 |
st.title("Simple Chat App using DistilBert Model (HuggingFace & Streamlit)")
|
45 |
|