diff --git a/development/index.html b/development/index.html
index e667447..7439887 100644
--- a/development/index.html
+++ b/development/index.html
@@ -11,6 +11,7 @@
         window.API_SERVICE_ENVIRONMENT = { port: 5000 };
         window.APP_STATE = {
             analysisResults: [],
+            selectedTextIds: [],
             error: null,
             loading: false
         };
diff --git a/frontend/MainContainer.tsx b/frontend/MainContainer.tsx
index 9cecc88..5f92678 100644
--- a/frontend/MainContainer.tsx
+++ b/frontend/MainContainer.tsx
@@ -5,7 +5,8 @@ import {
   requestGenerateText,
   setGeneratedTextResults,
   generateTextFailed,
-  generateText
+  generateText,
+  selectText
 } from './actions';
 import TextGenerationControl from "./components/TextGenerationControl";
 import TextGenerationResults from "./components/TextGenerationResults";
@@ -16,36 +17,25 @@ type ContainerProps = {
   analysisResults: AnalyzedText[],
   error: any,
   loading: any,
+  selectedTextIds: Array<number>,
   requestGenerateText: Function,
   setGeneratedTextResults: Function,
   generateTextFailed: any,
   generateText: any
+  selectText: Function
 }
 
 const Container: React.FunctionComponent<ContainerProps> = ({
   analysisResults,
   error,
   loading,
+  selectedTextIds,
   requestGenerateText,
   setGeneratedTextResults,
   generateTextFailed,
-  generateText
+  generateText,
+  selectText
   }) => {
-    const [selectedTextIds, setSelectedTextIds] = useState(new Array<number>());
-    console.log("Container state: ", selectedTextIds)
-
-    const onSelectTextId = (id: number) => {
-      if (selectedTextIds.includes(id)) {
-        console.log(`Deselecting ${id}`)
-        // Deselect
-        setSelectedTextIds(selectedTextIds.filter(item => item != id));
-      } else {
-        console.log(`Selecting ${id}`)
-        // Add to selected ids
-        setSelectedTextIds([...selectedTextIds, id]);
-      }
-    }
-
     const getDetailedAnalysisComponent = () => {
       if (selectedTextIds.length == 0) {
         return (<React.Fragment />);
@@ -65,7 +55,7 @@ const Container: React.FunctionComponent<ContainerProps> = ({
             analysisResults={analysisResults}
             loading={loading}
             error={error}
-            onSelectTextId={onSelectTextId}
+            onSelectTextId={selectText}
             selectedTextIds={selectedTextIds} />
           {getDetailedAnalysisComponent()}
         </div>
@@ -78,6 +68,7 @@ function mapStateToProps (state) {
     analysisResults: state.analysisResults,
     error: state.error,
     loading: state.loading,
+    selectedTextIds: state.selectedTextIds,
   };
 }
 
@@ -86,7 +77,8 @@ function mapDispatchToProps (dispatch) {
     requestGenerateText: requestGenerateText,
     setGeneratedTextResults: setGeneratedTextResults,
     generateTextFailed: generateTextFailed,
-    generateText: generateText
+    generateText: generateText,
+    selectText: selectText
   }, dispatch);
  }
 
diff --git a/frontend/actions.ts b/frontend/actions.ts
index 22e9237..48565e4 100644
--- a/frontend/actions.ts
+++ b/frontend/actions.ts
@@ -30,6 +30,12 @@ function generateText(config) {
     }
 }
 
+function selectText(id: number) {
+  return {
+    type: "SELECT_TEXT",
+    id: id
+  }
+}
 
 
 export {
@@ -37,4 +43,5 @@ export {
   setGeneratedTextResults,
   generateTextFailed,
   generateText,
+  selectText,
 };
diff --git a/frontend/app.css b/frontend/app.css
index 09e12a3..fa85492 100644
--- a/frontend/app.css
+++ b/frontend/app.css
@@ -24,4 +24,24 @@ h1 {
 .plot-spider {
   display: flex;
   justify-content: center;
+}
+
+.light-gray-text {
+  color: #8A8886;
+}
+
+.text-gen-form-grid {
+  display: grid;
+  grid-template-columns: 1fr 2fr 2fr 3fr;
+}
+
+.stacked-text-field-grid {
+  display: grid;
+  grid-template-rows: 1fr 1fr;
+  grid-template-columns: 1fr 2fr;
+  align-items: center;
+}
+
+.ms-Checkbox-label {
+  font-weight: 600;
 }
\ No newline at end of file
diff --git a/frontend/components/PerspectiveScoresBarChart.tsx b/frontend/components/PerspectiveScoresBarChart.tsx
index 27e5fc1..1eadd10 100644
--- a/frontend/components/PerspectiveScoresBarChart.tsx
+++ b/frontend/components/PerspectiveScoresBarChart.tsx
@@ -11,6 +11,7 @@ type PerspectiveScoresBarChartState = {
 
 type PerspectiveScoresBarChartProps = {
   id: number,
+  width: number,
   scores: any,
   defaultSelectedScore: string
 }
@@ -52,7 +53,7 @@ class PerspectiveScoresBarChart extends Component<PerspectiveScoresBarChartProps
 
     var margin = { top: 5, right: 15, bottom: 50, left: 55 }
     var h = 111 - margin.top - margin.bottom
-    var w = 200;
+    var w = this.props.width;
 
     // SVG
     d3.select(`#${this.domID}`).remove();
@@ -66,7 +67,7 @@ class PerspectiveScoresBarChart extends Component<PerspectiveScoresBarChartProps
         yScale = d3.scaleLinear().range([h, 0]);
 
     var g = svg.append("g")
-               .attr("transform", "translate(" + 55 + "," + margin.top + ")");
+               .attr("transform", "translate(0," + margin.top + ")");
 
     xScale.domain(this.props.scores.map(function(d) { return d.score; }));
     yScale.domain([0.0, 1.0]);
@@ -78,8 +79,8 @@ class PerspectiveScoresBarChart extends Component<PerspectiveScoresBarChartProps
      .attr("transform", "translate(0," + h + ")")
      .append("text")
      .attr("y", 30)
-     .attr("x", (w + margin.left)/2)
-     .attr("text-anchor", "end")
+     .attr("x", w/2)
+     .attr("text-anchor", "middle")
      .attr("fill", "black")
      .attr("font-size", "14px")
      .text(selectedScoreMessage);
diff --git a/frontend/components/TextGenerationControl.tsx b/frontend/components/TextGenerationControl.tsx
index 53b4e2d..7557f25 100644
--- a/frontend/components/TextGenerationControl.tsx
+++ b/frontend/components/TextGenerationControl.tsx
@@ -2,6 +2,7 @@ import React, { FunctionComponent } from "react";
 import ReactDOM from "react-dom";
 import { Fabric } from "office-ui-fabric-react/lib/Fabric";
 import { TextField } from "office-ui-fabric-react/lib/TextField";
+import { MaskedTextField } from "office-ui-fabric-react/lib/TextField";
 import { PrimaryButton } from "office-ui-fabric-react/lib/Button";
 import { Checkbox, Stack, Slider } from "office-ui-fabric-react";
 import { FiChevronUp, FiChevronDown } from "react-icons/fi";
@@ -127,12 +128,15 @@ class TextGenerationControl extends React.Component<
   }
 
   render() {
+    const textFieldStyle = {root: {width: "110px", display: "inline-block"}};
     if (this.state.isFolded) {
       return (
         <div
           style={{
               backgroundColor: "white",
-              padding: "24px",
+              paddingLeft: "24px",
+              paddingRight: "24px",
+              paddingTop: "24px",
               border: "1px solid #E7E7E7",
               borderRadius: "4px",
               boxShadow: "0px 2px 12px rgba(0, 0, 0, 0.12)",
@@ -149,7 +153,7 @@ class TextGenerationControl extends React.Component<
               />
           </div>
           <button className="flex items-center justify-center w-full p-1" onClick={() => this.toggleFold()}>
-              <FiChevronDown />
+              <FiChevronDown size={25} className="light-gray-text" />
           </button>
         </div>
       )
@@ -158,7 +162,9 @@ class TextGenerationControl extends React.Component<
         <div
           style={{
               backgroundColor: "white",
-              padding: "24px",
+              paddingLeft: "24px",
+              paddingRight: "24px",
+              paddingTop: "24px",
               border: "1px solid #E7E7E7",
               borderRadius: "4px",
               boxShadow: "0px 2px 12px rgba(0, 0, 0, 0.12)",
@@ -180,56 +186,88 @@ class TextGenerationControl extends React.Component<
               onClick={this.generateTextClickHandler}
             />
           </div>
-          <div className="grid grid-rows-2 gap-4 grid-flow-col p-5" style={{backgroundColor: "#F3F6FD"}}>
-            <Checkbox
-              label="Sample"
-              onChange={this.onSampleChange}
-              checked={this.state.doSample}
-            />
-            <Checkbox
-              label="Early Stopping"
-              onChange={this.onEarlyStoppingChange}
-              checked={this.state.earlyStopping}
-            />
-            <TextField
-              label="Max Length"
-              onChange={this.onMaxLengthChange}
-              value={this.state.maxLength.toString()}
-            />
-            <TextField
-              label="Min Length"
-              onChange={this.onMinLengthChange}
-              value={this.state.minLength.toString()}
-            />
-            <TextField
-              label="Top k"
-              onChange={this.onTopKChange}
-              value={this.state.topK.toString()}
-            />
-            <TextField
-              label="Num. Beams"
-              onChange={this.onNumBeamsChange}
-              value={this.state.numBeams.toString()}
-            />
-            <Slider
-              label="Temperature"
-              min={0.0}
-              max={1.0}
-              step={0.1}
-              onChange={this.onTemperatureChange}
-              value={this.state.temperature}
-            />
-            <Slider
-              label="Top p"
-              min={0.0}
-              max={1.0}
-              step={0.1}
-              onChange={this.onTopPChange}
-              value={this.state.topP}
-            />
+          <div className="text-gen-form-grid p-4" style={{backgroundColor: "#F3F6FD"}}>
+            <div className="grid grid-rows-2 items-center">
+              <Checkbox
+                label="Sample"
+                onChange={this.onSampleChange}
+                checked={this.state.doSample}
+              />
+              <Checkbox
+                label="Early Stopping"
+                onChange={this.onEarlyStoppingChange}
+                checked={this.state.earlyStopping}
+              />
+            </div>
+            <div className="stacked-text-field-grid">
+              <span className="mr-3 font-semibold text-right">Max Length</span>
+              <div>
+                <MaskedTextField
+                  onChange={this.onMaxLengthChange}
+                  value={this.state.maxLength.toString()}
+                  mask="999"
+                  maskChar=""
+                  styles={textFieldStyle}
+                />
+                <span className="ml-3 light-gray-text">[1-100]</span>
+              </div>
+              <span className="mr-3 font-semibold text-right">Min Length</span>
+              <div>
+                <MaskedTextField
+                  onChange={this.onMinLengthChange}
+                  value={this.state.minLength.toString()}
+                  mask="999"
+                  maskChar=""
+                  styles={textFieldStyle}
+                />
+                <span className="ml-3 light-gray-text">[1-100]</span>
+              </div>
+            </div>
+            <div className="stacked-text-field-grid">
+              <span className="mr-3 font-semibold text-right">Num. Beams</span>
+              <div>
+                <MaskedTextField
+                  onChange={this.onNumBeamsChange}
+                  value={this.state.numBeams.toString()}
+                  mask="999"
+                  maskChar=""
+                  styles={textFieldStyle}
+                />
+                <span className="ml-3 light-gray-text">[1-100]</span>
+              </div>
+              <span className="mr-3 font-semibold text-right">Top K</span>
+              <div>
+                <MaskedTextField
+                  onChange={this.onTopKChange}
+                  value={this.state.topK.toString()}
+                  mask="999"
+                  maskChar=""
+                  styles={textFieldStyle}
+                />
+                <div/>
+              </div>
+            </div>
+            <div className="grid grid-rows-2">
+              <Slider
+                label="Temperature"
+                min={0.0}
+                max={1.0}
+                step={0.1}
+                onChange={this.onTemperatureChange}
+                value={this.state.temperature}
+              />
+              <Slider
+                label="Top P"
+                min={0.0}
+                max={1.0}
+                step={0.1}
+                onChange={this.onTopPChange}
+                value={this.state.topP}
+              />
+            </div>
           </div>
           <button className="flex items-center justify-center w-full p-1" onClick={() => this.toggleFold()}>
-              <FiChevronUp />
+              <FiChevronUp size={25} className="light-gray-text" />
           </button>
         </div>
       );
diff --git a/frontend/components/TextGenerationResults.tsx b/frontend/components/TextGenerationResults.tsx
index 72d8247..217697a 100644
--- a/frontend/components/TextGenerationResults.tsx
+++ b/frontend/components/TextGenerationResults.tsx
@@ -1,11 +1,6 @@
 import React from "react";
-import * as _ from "lodash"
+import { ComboBox, IComboBoxOption, Stack } from "office-ui-fabric-react"
 import AnalyzedText from "../models/AnalyzedText";
-import ReactDOM from "react-dom";
-import { Fabric } from "office-ui-fabric-react/lib/Fabric";
-import { TextField } from "office-ui-fabric-react/lib/TextField";
-import { DefaultButton } from "office-ui-fabric-react/lib/Button";
-import { List } from "office-ui-fabric-react/lib/List";
 import PerspectiveScoresBarChart from "./PerspectiveScoresBarChart";
 
 type TextGenerationResultsProps = {
@@ -54,16 +49,53 @@ class TextGenerationResults extends React.Component<TextGenerationResultsProps,
       return resultScores;
     }
 
+    const sortOptions: IComboBoxOption[] = [
+      { key: 'toxicity', text: 'Toxicity' },
+      { key: 'severe', text: 'Severe Toxicity' },
+      { key: 'identity', text: 'Identity Attack' },
+      { key: 'insult', text: 'Insult' },
+      { key: 'profanity', text: 'Profanity' },
+      { key: 'threat', text: 'Threat' },
+      { key: 'sexually', text: 'Sexually Explicit' },
+      { key: 'flirtation', text: 'Flirtation' },
+    ];
+
+    const scoreOptions: IComboBoxOption[] = [
+      { key: 'all', text: 'All Scores' },
+      { key: 'toxicity', text: 'Toxicity' },
+      { key: 'severe', text: 'Severe Toxicity' },
+      { key: 'identity', text: 'Identity Attack' },
+      { key: 'insult', text: 'Insult' },
+      { key: 'profanity', text: 'Profanity' },
+      { key: 'threat', text: 'Threat' },
+      { key: 'sexually', text: 'Sexually Explicit' },
+      { key: 'flirtation', text: 'Flirtation' },
+    ];
+
+    const chartWidth = 140;
+
     return (
       <div>
-        <h1 className="mb-4">Generated Results</h1>
+        <h1>Generated Results</h1>
+        <Stack horizontal horizontalAlign="end" tokens={{childrenGap: 20}} styles={{root: {marginBottom: "16px"}}}>
+          <ComboBox
+            label="Sort by"
+            options={sortOptions}
+            styles={{root: {maxWidth: 160}}}
+          />
+          <ComboBox
+            label="Show score"
+            options={scoreOptions}
+            styles={{root: {maxWidth: 160}}}
+          />
+        </Stack>
         {this.props.analysisResults.map((item) => 
           <div className="flex items-stretch mb-4" style={{minHeight: "110px", border: this.props.selectedTextIds.includes(item.id) ? "solid #167DF5 2px" : "none"}}>
             <div className="flex items-center flex-none px-1" style={{backgroundColor: "#8894B1", color: "white", fontSize: "18px", lineHeight: "20px"}}>
               {item.id}
             </div>
-            <div className="flex-none" style={{width: "140px", border: "solid black 1px"}}>
-              <PerspectiveScoresBarChart id={item.id} scores={getResultScoreObj(item)} defaultSelectedScore="TOXICITY" />
+            <div className="flex-none" style={{width: `${chartWidth}px`, border: "solid black 1px"}}>
+              <PerspectiveScoresBarChart id={item.id} width={chartWidth} scores={getResultScoreObj(item)} defaultSelectedScore="TOXICITY" />
             </div>
             <button 
               className="flex items-center flex-auto text-left px-4"
diff --git a/frontend/reducers.ts b/frontend/reducers.ts
index 308412f..a595e05 100644
--- a/frontend/reducers.ts
+++ b/frontend/reducers.ts
@@ -1,9 +1,10 @@
-import AnalyzedText from "./models/AnalyzedText"
+import AnalyzedText from "./models/AnalyzedText";
 
 const rawInitialState = {
    analysisResults: [],
+   selectedTextIds: [],
    error: null,
-   loading: false
+   loading: false,
 }
 
 const initialState = window["APP_STATE"] || rawInitialState;
@@ -20,6 +21,16 @@ function rootReducer(state = initialState, action) {
         case "SET_GENERATED_TEXT_RESULTS":
           const analysisResults = action.text_generation_results?.map((result) => new AnalyzedText(result));
           return Object.assign({}, state, {loading: false, analysisResults: analysisResults, error: null});
+        
+        case "SELECT_TEXT":
+          const id = action.id;
+          let textIds = state.selectedTextIds;
+          if (state.selectedTextIds.includes(id)) {
+            textIds = textIds.filter(item => item != id);
+          } else {
+            textIds = [...textIds, id];
+          }
+          return Object.assign({}, state, {selectedTextIds: textIds});
 
         default:
             return state
diff --git a/transformerviz/assets/index.html b/transformerviz/assets/index.html
index a5b9d5b..0dc343e 100644
--- a/transformerviz/assets/index.html
+++ b/transformerviz/assets/index.html
@@ -13,6 +13,7 @@
 
         window.APP_STATE = {
             analysisResults: [],
+            selectedTextIds: [],
             error: null,
             loading: false
         };