From 0066dee420324e009d590279af6c0f9439afa7d1 Mon Sep 17 00:00:00 2001
From: Justus Schock <12886177+justusschock@users.noreply.github.com>
Date: Wed, 8 Feb 2023 16:01:49 +0100
Subject: [PATCH] next-version e2e (#12)

---
 tests/test_app.py    |  5 -----
 tests/test_module.py | 19 +++++--------------
 2 files changed, 5 insertions(+), 19 deletions(-)

diff --git a/tests/test_app.py b/tests/test_app.py
index 0340194..e8aa567 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -1,8 +1,5 @@
 import logging
-import os
 import sys
-from contextlib import redirect_stdout, redirect_stderr
-import io
 from typing import Union
 
 import lightning
@@ -15,13 +12,11 @@
 from lit_llms.tensorboard import (
     MultiNodeLightningTrainerWithTensorboard,
 )
-from lightning.app.runners import MultiProcessRuntime
 
 from transformers import T5ForConditionalGeneration, T5TokenizerFast as T5Tokenizer
 
 from lai_tldr import TLDRLightningModule
 from tests.test_module import BoringModel
-from lightning.app.testing import LightningTestApp
 
 class DummyTLDR(TLDR):
     boring_model=True
diff --git a/tests/test_module.py b/tests/test_module.py
index 83b0876..d82332d 100644
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -1,8 +1,7 @@
 import os
 import string
-from collections import namedtuple
+from argparse import Namespace
 from random import choice
-from time import sleep
 
 import lightning as L
 import pandas as pd
@@ -12,8 +11,6 @@
 from lai_tldr import TLDRDataModule
 from lai_tldr.module import TLDRLightningModule
 
-return_type = namedtuple("return_type", ("loss", "logits"))
-
 
 class BoringModel(torch.nn.Module):
     def __init__(self, target_seq_length: int, vocab_size: int, embed_size: int = 10):
@@ -23,18 +20,12 @@ def __init__(self, target_seq_length: int, vocab_size: int, embed_size: int = 10
         self.seq_length = target_seq_length
         self.vocab_size = vocab_size
 
-    def forward(self, input_ids, labels=None, **kwargs):
+    def forward(self, input_ids, **kwargs):
         # mimic source_seq_length -> target_seq_length by truncation
         logits = self.layer2(self.wte(input_ids))[:, : self.seq_length, :]
-        if labels is None:
-            loss = None
-        else:
-            loss = torch.nn.functional.cross_entropy(
-                logits.contiguous().view(-1, logits.size(-1)),
-                labels.view(-1),
-                ignore_index=-100,
-            )
-        return return_type(loss, logits)
+
+        loss = logits.sum()
+        return Namespace(loss=loss, logits=logits)
 
     def generate(self, input_ids, num_beams: int = 1, **kwargs):
         return [self(input_ids).logits.argmax(-1)[0].tolist() for _ in range(num_beams)]