Scraping with some delay

[Truman]

from selenium import webdriver
from selenium.webdriver.support import ui

driver = webdriver.Firefox()
driver.get("http://the-internet.herokuapp.com/dynamic_loading/2")
button = driver.find_element_by_xpath("//*/div[@id='start']/button")

button.click()
print("clicked")

wait = ui.WebDriver.Wait(driver, 10)
wait.until(lambda driver: driver.find_element_by_xpath("//*/div[@id='finish']"))

finish_element=driver.find_element_by_xpath("//*/div[@id='finish']/h4")
print(finish_element.text)

Error:
Traceback (most recent call last):
  File "C:\Python36\kodovi\click_button.py", line 11, in <module>
    wait = ui.WebDriver.Wait(driver, 10)
AttributeError: module 'selenium.webdriver.support.ui' has no attribute 'WebDriv
er'

Don't understand why there isn't a webdriver attribute. What should I put instead to do a scraping after the page is in status 'finish'? How to create object 'wait'?