Update app.py

app.py

@@ -1,6 +1,6 @@
 import re
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import pipeline
 
 # Initialize the chat history
 history = []
@@ -8,38 +8,30 @@ history = []
 def clean_text(text):
     return re.sub('[^a-zA-Z\s]', '', text).strip()
 
-tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
-model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/DialoGPT-small").half().cuda()
+# Load DistilBert model
+model = pipeline("question-answering", model="distilbert-base-cased")
 
 def generate_response(user_input):
+    # Add user input to history
     history.append((user_input, ""))
 
     if not history:
         return ""
     
     last_user_message = history[-1][0]
-    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, decoded_output)
-    return f"AI: {decoded_output}".capitalize()
+    user_input = clean_text(last_user_message)
+
+    if len(user_input) > 0:
+        # Modify the context value below with appropriate content
+        #Replace this line with useful information for generating proper responses.
+        context = """The Phoenix pay system is a payroll processing system for Canadian federal government employees, provided by IBM in June 2011 using PeopleSoft software, and run by Public Services and Procurement Canada. The Public Service Pay Centre is located in Miramichi, New Brunswick. It was first introduced in 2009 as part of Prime Minister Stephen Harper's Transformation of Pay Administration Initiative, intended to replace Canada's 40-year old system with a new, cost-saving "automated, off-the-shelf commercial system."
+By July 2018, Phoenix has caused pay problems to close to 80 percent of the federal government's 290,000 public servants through underpayments, over-payments, and non-payments.[1] The Standing Senate Committee on National Finance, chaired by Senator Percy Mockler, sought to examine the causes for the failure, holding "eight meetings with 28 witnesses, including the Auditor General of Canada, union representatives, departments and agencies, officials from IBM, the Minister of Public Services and Procurement and the Clerk of the Privy Council"[1] and paid a visit to the Miramichi pay system location during their investigation. Their report, "The Phoenix Pay Problem: Working Towards a Solution" on July 31, 2018, in which they called Phoenix a failure and an "international embarrassment".[1] Instead of saving $70 million a year as planned, the report said that the cost to taxpayers to fix Phoenix's problems could reach a total of $2.2 billion by 2023. The Office of the Auditor General of Canada also performed an independent audit, and published a report in 2018 that concluded that the Phoenix project “was a incomprehensible failure of project management and oversight”,[2] and that Phoenix Executives did not heed warnings from the Miramichi Pay Centre, costing the federal government hundreds of millions of dollars, and had a negative financial impact on tens of thousands of its employees. """
+        
+        result = model(question=user_input, context=context)
+        answer = result['answer']
+        history[-1] = (last_user_message, answer)
+
+        return f"AI: {answer}"
 
 st.title("Simple Chat App using DistilBert Model (HuggingFace & Streamlit)")