Update app.py

app.py

@@ -1,40 +1,53 @@
-import torch
-from transformers import DistilBertTokenizerFast, DistilBertForQuestionAnswering
-import json
+import re
 import streamlit as st
 from transformers import pipeline
 
-model_name = "distilbert-base-cased"
-tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
-model = DistilBertForQuestionAnswering.from_pretrained(model_name)
+# Initialize the chat history
+history = []
 
-def get_user_input():
-    return st.chat_input("You:", "")
+def clean_text(text):
+    return re.sub('[^a-zA-Z\s]', '', text).strip()
 
-def generate_response(prompt):
-    try:
-        response = model(prompt, max_length=50, do_sample=True)
-        return response[0]['generated_text'].strip()
-    except Exception as e:
-        print(f"Error while generating response: {str(e)}")
-        return "Sorry, I cannot respond right now."
+# Load DistilBert model
+model = pipeline("question-answering", model="distilbert-base-cased")
 
-if __name__ == "__main__":
-    st.set_page_config(page_title="AI Chat App", page_icon=":robot_face:", layout="wide")
-    
-    st.write("# Welcome to the AI Chat App! :wave:")
-    st.markdown("<style>.container { width:100% !important; }</style>", unsafe_allow_html=True)
+def generate_response(user_input):
+    # Add user input to history
+    history.append((user_input, ""))
+
+    if not history:
+        return ""
     
-    user_input = ""
-    while True:
-        # Get user input
-        user_input = get_user_input()
-        
-        if not user_input:
-            continue
-        
-        # Generate response
-        response = generate_response(user_input)
-        
-        # Display response
-        st.markdown(f"<b>AI:</b> {response}", unsafe_allow_html=True)
+    last_user_message = history[-1][0]
+    user_input = clean_text(last_user_message)
+
+    if len(user_input) > 0:
+        result = model(question=user_input, context="Placeholder text")
+        answer = result['answer']
+        history[-1] = (last_user_message, answer)
+
+        return f"AI: {answer}"
+
+st.title("Simple Chat App using DistilBert Model (HuggingFace & Streamlit)")
+
+for i in range(len(history)):
+    message = history[i][0]
+    response = history[i][1]
+
+    if i % 2 == 0:
+        col1, col2 = st.beta_columns([0.8, 0.2])
+        with col1:
+            st.markdown(f">> {message}")
+        with col2:
+            st.write("")
+    else:
+        col1, col2 = st.beta_columns([0.8, 0.2])
+        with col1:
+            st.markdown(f"   {response}")
+        with col2:
+            st.button("Clear")
+
+new_message = st.text_area("Type something...")
+if st.button("Submit"):
+    generated_response = generate_response(new_message)
+    st.markdown(generated_response)