update

pages/summarization_playground.py

@@ -58,7 +58,7 @@ def generate_answer(sources, model_name, prompt):
 
     content = [prompt + '\n{' + sources + '}\n\nsummary:']
 
-    answer = model[model_name].gen(content, streaming=False)
+    answer = model[model_name].gen(content)
 
     return answer
 

utils/model.py

@@ -53,56 +53,59 @@ class Model(torch.nn.Module):
         return self.tokenizer
     
     def return_model(self):
-        return self.pipeline
-
-    def gen(self, content_list, temp=0.001, max_length=500, streaming=False):
+        return self.model
+    
+    def streaming(self, content_list, temp=0.001, max_length=500):
         # Convert list of texts to input IDs
         input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
 
-        if streaming:
-            # Set up the initial generation parameters
-            gen_kwargs = {
-                "input_ids": input_ids,
-                "do_sample": True,
-                "temperature": temp,
-                "eos_token_id": self.tokenizer.eos_token_id,
-                "max_new_tokens": 1,  # Generate one token at a time
-                "return_dict_in_generate": True,
-                "output_scores": True
-            }
-
-            # Generate and yield tokens one by one
-            generated_tokens = 0
-            batch_size = input_ids.shape[0]
-            active_sequences = torch.arange(batch_size)
-
-            while generated_tokens < max_length and len(active_sequences) > 0:
-                with torch.no_grad():
-                    output = self.model.generate(**gen_kwargs)
-                
-                next_tokens = output.sequences[:, -1].unsqueeze(-1)
-                
-                # Yield the newly generated tokens for each sequence in the batch
-                for i, token in zip(active_sequences, next_tokens):
-                    yield i, self.tokenizer.decode(token[0], skip_special_tokens=True)
-
-                # Update input_ids for the next iteration
-                gen_kwargs["input_ids"] = torch.cat([gen_kwargs["input_ids"], next_tokens], dim=-1)
-                generated_tokens += 1
-
-                # Check for completed sequences
-                completed = (next_tokens.squeeze(-1) == self.tokenizer.eos_token_id).nonzero().squeeze(-1)
-                active_sequences = torch.tensor([i for i in active_sequences if i not in completed])
-                if len(active_sequences) > 0:
-                    gen_kwargs["input_ids"] = gen_kwargs["input_ids"][active_sequences]
+        # Set up the initial generation parameters
+        gen_kwargs = {
+            "input_ids": input_ids,
+            "do_sample": True,
+            "temperature": temp,
+            "eos_token_id": self.tokenizer.eos_token_id,
+            "max_new_tokens": 1,  # Generate one token at a time
+            "return_dict_in_generate": True,
+            "output_scores": True
+        }
+
+        # Generate and yield tokens one by one
+        generated_tokens = 0
+        batch_size = input_ids.shape[0]
+        active_sequences = torch.arange(batch_size)
+
+        while generated_tokens < max_length and len(active_sequences) > 0:
+            with torch.no_grad():
+                output = self.model.generate(**gen_kwargs)
+            
+            next_tokens = output.sequences[:, -1].unsqueeze(-1)
+            
+            # Yield the newly generated tokens for each sequence in the batch
+            for i, token in zip(active_sequences, next_tokens):
+                yield i, self.tokenizer.decode(token[0], skip_special_tokens=True)
+
+            # Update input_ids for the next iteration
+            gen_kwargs["input_ids"] = torch.cat([gen_kwargs["input_ids"], next_tokens], dim=-1)
+            generated_tokens += 1
+
+            # Check for completed sequences
+            completed = (next_tokens.squeeze(-1) == self.tokenizer.eos_token_id).nonzero().squeeze(-1)
+            active_sequences = torch.tensor([i for i in active_sequences if i not in completed])
+            if len(active_sequences) > 0:
+                gen_kwargs["input_ids"] = gen_kwargs["input_ids"][active_sequences]
+
+
+    def gen(self, content_list, temp=0.001, max_length=500):
+        # Convert list of texts to input IDs
+        input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
 
-        else:
-            # Non-streaming generation (unchanged)
-            outputs = self.model.generate(
-                input_ids,
-                max_new_tokens=max_length,
-                do_sample=True,
-                temperature=temp,
-                eos_token_id=self.tokenizer.eos_token_id,
-            )
-            return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        # Non-streaming generation (unchanged)
+        outputs = self.model.generate(
+            input_ids,
+            max_new_tokens=max_length,
+            do_sample=True,
+            temperature=temp,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)

utils/multiple_stream.py

@@ -20,7 +20,7 @@ def stream_data(content_list, model):
     outputs = ["" for _ in content_list]
 
     # Use the gen method to handle batch generation
-    generator = model.gen(content_list, streaming=True)
+    generator = model.streaming(content_list)
     
     while True:
         updated = False