[RichUI] Better Dataset Generation Display

[benjaminye]

Is your feature request related to a problem? Please describe.

• Dataset creation table display always display all columns of dataset, instead of ones needed by prompt and prompt_stub
• Dataset creation table display highlighting uses string matching, leading to weird outputs when there are overlaps

Describe the solution you'd like

• Fix these issues!