Update app.py

app.py

@@ -2,53 +2,39 @@ import torch
 from transformers import DistilBertTokenizerFast, DistilBertForQuestionAnswering
 import json
 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)
 
-def handle_query_change(_, query):
-    update_displayed_query(query)
+def get_user_input():
+    return st.chat_input("You:", "")
 
-def format_response(start_index, end_index, raw_answer):
-    answer_tokens = tokenizer.convert_tokens_to_string([tokenizer.convert_ids_to_tokens(i)[0] for i in range(start_index, end_index+1)])
-    return {'answer': answer_tokens.strip(), 'score': None}
-
-def get_answers(question, context):
-    inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
-    start_scores, end_scores = model(**inputs).values()
-    start_index = torch.argmax(start_scores)
-    end_index = torch.argmax(end_scores) + 1
-    formatted_answer = format_response(start_index, end_index - 1, context[start_index:end_index].tolist())
-    return formatted_answer
-
-def update_displayed_query(query):
-    st.session_state['current_query'] = query
-
-def interactive():
-    #print("Hi! I am a simple AI chatbot built using Hugging Face and Streamlit.")
-    query = ""
-    while query != st.session_state.get('current_query', ''):
-        if ('current_query' not in st.session_state) or (st.session_state['current_query'] != query):
-            st.session_state['current_query'] = query
-            st.text_input("Ask me something or type 'quit' to exit:", "", key='my_input', on_change=handle_query_change, help="Enter your question here...") # added the help parameter to show the hint
-            st.echo(False) # hide the default output of the text input widget
-
-        if ('shown_query' not in st.session_state) or (st.session_state['shown_query'] != query):
-            st.session_state['shown_query'] = query
-            if len(query) > 0 and query != 'quit':
-                try:
-                    context = "The capital of France is Paris."
-                    response = get_answers(query, context)
-                    st.write(json.dumps(response))
-                except Exception as e:
-                    st.write(f"Error occurred: {str(e)}")
-
-        query = st.session_state.get('current_query', '')
+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."
 
 if __name__ == "__main__":
-    st.set_page_config(layout="wide")
-    st.title('AI Chatbot Built Using Hugging Face and Streamlit')
-    st.subheader('Welcome to Our Intelligent Conversational Agent!')
-    st.write('Please enter your question below or type "quit" to close the session.')
-    interactive()
+    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)
+    
+    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)