From 910b850922293c7d7f09e11c1e776e491d63b05f Mon Sep 17 00:00:00 2001
From: Xavier Fernandes <fernandes.praphat@gmail.com>
Date: Wed, 26 May 2021 17:21:53 -0700
Subject: [PATCH] Some tweaks ahead of the demo

Fix the random seed so that we get consistent results that are reproducable.
Hide the ascending / descending buttons if sort is set to none.
After generating a bunch of text, reset the dropdowns state.
---
 frontend/components/TextGenerationResults.tsx | 31 +++++++++++++------
 transformerviz/helpers/utils.py               |  7 +++++
 transformerviz/server.py                      |  3 ++
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/frontend/components/TextGenerationResults.tsx b/frontend/components/TextGenerationResults.tsx
index 2e05b6a..aadfd37 100644
--- a/frontend/components/TextGenerationResults.tsx
+++ b/frontend/components/TextGenerationResults.tsx
@@ -34,7 +34,10 @@ class TextGenerationResults extends React.Component<TextGenerationResultsProps,
 
   componentWillReceiveProps(nextProps) {
     this.setState({
-      analysisResults: nextProps.analysisResults
+      analysisResults: nextProps.analysisResults,
+      highlightScoreLabel: "toxicity",
+      sortByLabel: "none",
+      sortDescending: true
     });
   }
 
@@ -47,11 +50,11 @@ class TextGenerationResults extends React.Component<TextGenerationResultsProps,
     } else if (sortDescending) {
       const sortedAnalysisResultsDescending = this.state.analysisResults.map((result) => result).sort((resultA, resultB) => {
         if (resultA[scoreLabel] < resultB[scoreLabel]) {
-          return -1;
+          return 1;
         } else if (resultA[scoreLabel] == resultB[scoreLabel]) {
           return 0;
         } else {
-          return 1;
+          return -1;
         }
       });
 
@@ -62,11 +65,11 @@ class TextGenerationResults extends React.Component<TextGenerationResultsProps,
     } else if (!sortDescending) {
       const sortedAnalysisResultsAscending = this.state.analysisResults.map((result) => result).sort((resultA, resultB) => {
         if (resultA[scoreLabel] < resultB[scoreLabel]) {
-          return 1;
+          return -1;
         } else if (resultA[scoreLabel] == resultB[scoreLabel]) {
           return 0;
         } else {
-          return -1;
+          return 1;
         }
       });
 
@@ -150,15 +153,25 @@ class TextGenerationResults extends React.Component<TextGenerationResultsProps,
 
     const chartWidth = 140;
 
+    const getAscendingDescendingButtons = () => {
+      if (this.state.sortByLabel != "none") {
+        return (
+          <div className="flex items-end">
+            <button className={getSortClass(this.state.sortDescending)} onClick={() => this.sortDescending()} ><IconSortDescending size={30} /></button>
+            <button className={getSortClass(!this.state.sortDescending)} onClick={() => this.sortAscending()} ><IconSortAscending size={30} /></button>
+          </div>
+        );
+      }
+
+      return (<React.Fragment />);
+    }
+
     return (
       <div className="page-column">
         <div className="scroll-pane">
         <h1>Generated Results</h1>
         <Stack horizontal horizontalAlign="end" tokens={{childrenGap: 20}} styles={{root: {marginBottom: "16px"}}}>
-          <div className="flex items-end">
-            <button className={getSortClass(this.state.sortDescending)} onClick={() => this.sortDescending()} ><IconSortDescending size={30} /></button>
-            <button className={getSortClass(!this.state.sortDescending)} onClick={() => this.sortAscending()} ><IconSortAscending size={30} /></button>
-          </div>
+          {getAscendingDescendingButtons()}
           <ComboBox
             label="Sort by"
             options={sortOptions}
diff --git a/transformerviz/helpers/utils.py b/transformerviz/helpers/utils.py
index 75dd892..be0cf96 100644
--- a/transformerviz/helpers/utils.py
+++ b/transformerviz/helpers/utils.py
@@ -3,6 +3,13 @@
     AutoModelForCausalLM)
 from torch.nn import functional as F
 from googleapiclient import discovery
+import torch
+
+
+def set_seed(seed):
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
 
 
 def generate_sentences_hf(tokenizer, model, context,
diff --git a/transformerviz/server.py b/transformerviz/server.py
index b776059..ecba7b5 100644
--- a/transformerviz/server.py
+++ b/transformerviz/server.py
@@ -11,6 +11,8 @@
 app = Flask(__name__)
 CORS(app)
 
+utils.set_seed(0)
+
 text_generators = {
     "gpt2": utils.TextGenerator(
         os.environ.get("GPT2_MODEL_VERSION", "gpt2"),
@@ -82,6 +84,7 @@ def generate_text():
     temperature = request_json.get("temperature")
     num_beams = request_json.get("num_beams")
     text_generator = text_generators.get(model_name)
+    utils.set_seed(0)
     generated_sentences = text_generator.generate_text(
         prompt,
         do_sample=do_sample,