Update app.py

app.py

@@ -57,18 +57,23 @@ collection = chroma_client.get_collection(name="F1-wiki", embedding_function=Cus
 
 question = get_text()
 
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
 if question:
-    response = collection.query(query_texts=question, include=['documents'], n_results=5)
+    with st.spinner("Generating answer... "):
+        response = collection.query(query_texts=question, include=['documents'], n_results=5)
 
-    context = " ".join(response['documents'][0])
+        context = " ".join(response['documents'][0])
 
-    input_text = template.replace("{context}", context).replace("{question}", question)
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+        input_text = template.replace("{context}", context).replace("{question}", question)
+        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+        attention_mask = (input_ids != tokenizer.pad_token_id).to(device)
 
-    output = model.generate(input_ids, max_new_tokens=200, early_stopping=True)
-    answer = tokenizer.decode(output[0], skip_special_tokens=True).split("ANSWER:")[1]
+        output = model.generate(input_ids, attention_mask=attention_mask, max_new_tokens=200, early_stopping=True)
+        answer = tokenizer.decode(output[0], skip_special_tokens=True).split("ANSWER:")[1].strip()
 
-    st.session_state.past.append(question)
-    st.session_state.generated.append(answer)
+        st.session_state.past.append(question)
+        st.session_state.generated.append(answer)
 
     st.write(answer)