Skip to content

Commit

Permalink
fix getitem
Browse files Browse the repository at this point in the history
  • Loading branch information
huyiwen committed Jul 10, 2024
1 parent 2806c25 commit 4ebeac5
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion utilization/model/model_utils/conversation.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,6 @@ def __init__(
model_evaluation_method: Optional[Literal["get_ppl", "get_prob", "generation", "user_defined"]] = None,
split: Optional[bool] = None,
is_normalized: bool = False,
**deprecated_kwargs
):
self.messages = messages if isinstance(messages, list) else []
self.num_turns = num_turns
Expand Down Expand Up @@ -454,5 +453,18 @@ def add_(
self.messages.extend(messages)
return self

def __iter__(self):
for message in self.messages:
yield message

def __getitem__(self, item):
return self.messages[item]

def __setitem__(self, key, value):
self.messages[key] = value

def __len__(self):
return len(self.messages)

def __repr__(self):
return "Conversation(\n" + pformat(self.messages) + ")"

0 comments on commit 4ebeac5

Please sign in to comment.