Update app.py

app.py

@@ -3,7 +3,8 @@ from transformers import DistilBertTokenizerFast, DistilBertForQuestionAnswering
 
 model_name = "distilbert-base-cased"
 tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
-model = DistilBertForQuestionAnswering.from_pretrained(model_name)
+# Set output_attentions and output_hidden_states to False to prevent weight initialization warnings
+model = DistilBertForQuestionAnswering.from_pretrained(model_name, output_attentions=False, output_hidden_states=False)
 
 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)])
@@ -22,16 +23,16 @@ def main():
     print("Type 'quit' to exit the program.")
     query = ""
     while True:
-        query = input("Your Question: ")
+        query = input("Your Question: ").strip()
         if query.lower() == "quit":
             break
-        if len(query) > 0:
+        elif len(query) > 0:
             context = "The capital of France is Paris."
             try:
                 response = get_answers(query, context)
-                print(f"Response: {response}")
+                print(f"\nResponse: {response}\n")
             except Exception as e:
-                print(f"Error occurred: {str(e)}")
+                print(f"\nError occurred: {str(e)}\n")
 
 if __name__ == "__main__":
     main()