id4thomas commited on
Commit
699d7ca
·
1 Parent(s): 126efa4

[edit] update readme

Browse files
Files changed (1) hide show
  1. README.md +10 -10
README.md CHANGED
@@ -35,21 +35,21 @@ The model is trained to predict in the following schema
35
  from enum import Enum
36
  from pydantic import BaseModel
37
 
38
- class RelationshipStatus(str, Enum):
39
  na = "na"
40
  low = "low"
41
  medium = "medium"
42
  high = "high"
43
 
44
  class EmotionLabel(BaseModel):
45
- joy: RelationshipStatus
46
- trust: RelationshipStatus
47
- fear: RelationshipStatus
48
- surprise: RelationshipStatus
49
- sadness: RelationshipStatus
50
- disgust: RelationshipStatus
51
- anger: RelationshipStatus
52
- anticipation: RelationshipStatus
53
 
54
  class EntryResult(BaseModel):
55
  emotion: EmotionLabel
@@ -81,7 +81,7 @@ input_text = tokenizer.apply_chat_template(
81
  add_generation_prompt=True,
82
  )
83
  prediction = generator(input_text)
84
- >>> EntryResult(emotion=EmotionLabel(joy=<RelationshipStatus.na: 'na'>, ...)
85
  ```
86
 
87
  Using endpoint loaded with vllm & OpenAI client package
 
35
  from enum import Enum
36
  from pydantic import BaseModel
37
 
38
+ class EmotionLevel(str, Enum):
39
  na = "na"
40
  low = "low"
41
  medium = "medium"
42
  high = "high"
43
 
44
  class EmotionLabel(BaseModel):
45
+ joy: EmotionLevel
46
+ trust: EmotionLevel
47
+ fear: EmotionLevel
48
+ surprise: EmotionLevel
49
+ sadness: EmotionLevel
50
+ disgust: EmotionLevel
51
+ anger: EmotionLevel
52
+ anticipation: EmotionLevel
53
 
54
  class EntryResult(BaseModel):
55
  emotion: EmotionLabel
 
81
  add_generation_prompt=True,
82
  )
83
  prediction = generator(input_text)
84
+ >>> EntryResult(emotion=EmotionLabel(joy=<EmotionLevel.na: 'na'>, ...)
85
  ```
86
 
87
  Using endpoint loaded with vllm & OpenAI client package