This exercise is from Bloch's book and can be found here. Bloch introduces equivalent variations of the axiom of choice where the one that will be proven is stated in terms of functions:
AC1 Let $I$ be a nonempty set, and let $\{A_i\}_{i \in I}$ be a family of nonempty sets indexed by I. Then there is a function $f: I \rightarrow \bigcup_{i \in I} A_i$ such that $f(i) \in A_i$ for all $i \in I$.
As you can read from the text from the above link, the author suggests that AC1 can be deduced from AC2 below
AC2 Let $I$ be a nonempty set, and let $\{A_i\}_{i \in I}$ be a family of nonempty, pairwise disjoint sets indexed by I. Then there is a function $f: I \rightarrow \bigcup_{i \in I} A_i$ such that $f(i) \in A_i$ for all $i \in I$.
Since the second definition implies the first, it is sufficient to prove the second version of AC, which will imply the first version.
We will prove that if every surjective function has a right inverse, then AC2 is true, from which it follows that AC1 is true.
Here is my attempt of a proof.
Suppose that every surjective function has a right inverse, and let $f: A \rightarrow B$ be such a surjective function with a right inverse $g: B \rightarrow A$ where $A$ and $B$ are nonempty sets. Since $f$ is surjective and $g$ is a right inverse of $f$ it follows that $g$ is injective (this fact can be proven by contradiction). Since $g$ is a function every $b \in B$ is defined, so that the image $g(\{b\}) \neq \varnothing$. Further, the injectivity of $g$ implies that $g(\{b_1\}) \cap g(\{b_2\}) = \varnothing$ for arbitrary $b_1, b_2 \in B$ such that $b_1 \neq b_2$. Observe also that $g(\{b\})$ contains only one element for each $b$. Finally, it is clear that $A = \bigcup_{b \in B} g(\{b\})$ since $g$ is injective,$f$ is defined on $A$, and $g$ is defined on $B$. If we now consider the family $\{g(\{b\})\}_{b \in B}$, it is clear that this collection contains nonempty, pairwise dsjoint sets. Because $A = \bigcup_{b \in B} g(\{b\})$, we have $g:B \rightarrow \bigcup_{b \in B} g(\{b\})$, which is the desired choice function in the second definition above. Therefore, the AC of the first defintiion follows.
$\blacksquare$
Any criticisms or advice on the above proof is most welcome. I am unsure about my reason for the set equality $A = \bigcup_{b \in B} g(\{b\})$, but it can be proven formally in the traditional manner in demonstrating that the two sets are subsets of each other, but I wanted to keep my reasons for their equality concise, so I extracted the main arguments from such a proof and put them in this proof for the AC.
Thanks!
Your proof does not prove $\sf AC2$ as stated. It does show that if $f\colon A\to B$ is a surjection, and $g$ is a right inverse, then $g$ is a choice function from a family of pairwise disjoint non-empty sets.
What you need to do is start with a family of pairwise disjoint non-empty sets, and then use that family of produce a surjection whose right inverse will give you (directly or indirectly) a choice function.