update

pages/summarization_playground.py

@@ -53,6 +53,7 @@ def generate_answer(sources, model_name, prompt):
 
 def process_input(input_text, model_selection, prompt):
     if input_text:
+        logging.info("Start generation")
         response = generate_answer(input_text, model_selection, prompt)
         return f"## Original Dialogue:\n\n{input_text}\n\n## Summarization:\n\n{response}"
     else:

utils/model.py

@@ -23,6 +23,7 @@ class Model(torch.nn.Module):
         
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.name = model_name
+        logging.info(f'start loading model {self.name}')
         self.pipeline = transformers.pipeline(
             "summarization",
             model=model_name,
@@ -31,7 +32,7 @@ class Model(torch.nn.Module):
             device_map="auto",
         )
 
-        logging.info(f'Load model {self.name}')
+        logging.info(f'Loaded model {self.name}')
         self.update()
 
     @classmethod
@@ -48,13 +49,24 @@ class Model(torch.nn.Module):
         return self.pipeline
 
     def gen(self, content, temp=0.1, max_length=500):
-        sequences = self.pipeline(
-            content,
-            max_new_tokens=max_length,
-            do_sample=True,
-            temperature=temp,
-            num_return_sequences=1,
-            eos_token_id=self.tokenizer.eos_token_id,
-        )
+        if self.name == "google-t5/t5-large":
+            sequences = self.pipeline(
+                content,
+                max_new_tokens=max_length,
+                do_sample=True,
+                temperature=temp,
+                num_return_sequences=1,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+        else:
+            sequences = self.pipeline(
+                content,
+                max_new_tokens=max_length,
+                do_sample=True,
+                temperature=temp,
+                num_return_sequences=1,
+                eos_token_id=self.tokenizer.eos_token_id,
+                return_full_text=False
+            )
         
         return sequences[-1]['summary_text']