cleaned up some code, added condition to yield