Update text_generation.py

text_generation.py

@@ -1,4 +1,3 @@
-
 import torch
 import torch.nn.functional as F
 from tqdm import trange
@@ -47,7 +46,8 @@ try:
     device
 except NameError:
     device = "cpu"
-if device.startswith("cuda"):
+
+if torch.device(device).type == "cuda":
     torch.backends.cudnn.benchmark = True
 
 MAX_GENERATION_LENGTH = 512
@@ -106,17 +106,38 @@ def _generate_sequence(model_call, context_tensor, generated, decode_fn, end_tok
 def sample_sequence(prompt, model, enc, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
     context_tokens = enc.encode(prompt)
     context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
-    return _generate_sequence(lambda ct, past: model(ct, past_key_values=past), context_tensor, list(context_tokens), lambda token: enc.decode([token]), lambda token: token == enc.encoder[END_OF_TEXT_TOKEN], temperature, top_k, top_p, repetition_penalty, max_length)
+    return _generate_sequence(
+        lambda ct, past: model(ct, past_key_values=past),
+        context_tensor,
+        list(context_tokens),
+        lambda token: enc.decode([token]),
+        lambda token: token == enc.encoder[END_OF_TEXT_TOKEN],
+        temperature, top_k, top_p, repetition_penalty, max_length
+    )
 
 def sample_sequence_codegen(prompt, model, tokenizer, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"):
     context_tokens = tokenizer.encode(prompt)
     context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device)
-    return _generate_sequence(lambda ct, past: model(input_ids=ct, past_key_values=past, labels=None), context_tensor, list(context_tokens), lambda token: tokenizer.decode([token]), lambda token: token == 50256, temperature, top_k, top_p, repetition_penalty, max_length)
+    return _generate_sequence(
+        lambda ct, past: model(input_ids=ct, past_key_values=past, labels=None),
+        context_tensor,
+        list(context_tokens),
+        lambda token: tokenizer.decode([token]),
+        lambda token: token == 50256,
+        temperature, top_k, top_p, repetition_penalty, max_length
+    )
 
 def summarize_text(text):
     if summarization_model and summarization_tokenizer:
         input_ids = summarization_tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=1024).to(device)
-        summary_ids = summarization_model.generate(input_ids, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
+        summary_ids = summarization_model.generate(
+            input_ids,
+            max_length=150,
+            min_length=40,
+            length_penalty=2.0,
+            num_beams=4,
+            early_stopping=True
+        )
         return summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return text[:300] + "..." if len(text) > 300 else text