Chris4K commited on
Commit
a49896f
·
verified ·
1 Parent(s): 029d41a

Update sentiment_analysis.py

Browse files
Files changed (1) hide show
  1. sentiment_analysis.py +5 -5
sentiment_analysis.py CHANGED
@@ -12,7 +12,7 @@ class SentimentAnalysisTool(Tool):
12
  "description": "The text to analyze for sentiment"
13
  }
14
  }
15
- output_type = "list"
16
 
17
  # Available sentiment analysis models
18
  models = {
@@ -38,13 +38,13 @@ class SentimentAnalysisTool(Tool):
38
  return self.predict(text)
39
 
40
  def _parse_output(self, output_json):
41
- """Parse model output into a list of (label, score) tuples."""
42
- list_pred = []
43
  for i in range(len(output_json[0])):
44
  label = output_json[0][i]['label']
45
  score = output_json[0][i]['score']
46
- list_pred.append((label, score))
47
- return list_pred
48
 
49
  def _get_classifier(self, model_id):
50
  """Get or create a classifier for the given model ID."""
 
12
  "description": "The text to analyze for sentiment"
13
  }
14
  }
15
+ output_type = "json"
16
 
17
  # Available sentiment analysis models
18
  models = {
 
38
  return self.predict(text)
39
 
40
  def _parse_output(self, output_json):
41
+ """Parse model output into a dictionary of scores by label."""
42
+ result = {}
43
  for i in range(len(output_json[0])):
44
  label = output_json[0][i]['label']
45
  score = output_json[0][i]['score']
46
+ result[label] = score
47
+ return result
48
 
49
  def _get_classifier(self, model_id):
50
  """Get or create a classifier for the given model ID."""