Skip to content

Commit

Permalink
Using AI tools to improve commenting of base G-retriever example
Browse files Browse the repository at this point in the history
  • Loading branch information
riship committed Dec 21, 2024
1 parent dfd16f6 commit 0b8d565
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions examples/llm/g_retriever.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,15 @@ def get_loss(model, batch, model_save_name: str) -> Tensor:
# For LLM models, compute the loss using the question, label, and desc inputs
return model(batch.question, batch.label, batch.desc)
else: # (GNN+LLM)
return model(batch.question, batch.x, batch.edge_index, batch.batch,
batch.label, batch.edge_attr, batch.desc)
return model(
batch.question,
batch.x, # node features
batch.edge_index, # edge indices
batch.batch, # batch indices
batch.edge_attr, # edge attributes
batch.label, # answers (labels)
batch.desc # description
)


def inference_step(model, batch, model_save_name):
Expand Down

0 comments on commit 0b8d565

Please sign in to comment.