Update humanize.py

humanize.py

@@ -1,93 +1,93 @@
-# import torch
-# from nltk import sent_tokenize
-# import nltk
-# from tqdm import tqdm
-# from transformers import T5ForConditionalGeneration, T5Tokenizer
+import torch
+from nltk import sent_tokenize
+import nltk
+from tqdm import tqdm
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 
-# nltk.download("punkt")
-# # autodetect the available device
-# GPU_IDX = 1  # which GPU to use
-# if torch.cuda.is_available():
-#     num_gpus = torch.cuda.device_count()
-#     print(f"Number of available GPUs: {num_gpus}")
-#     assert GPU_IDX < num_gpus, f"GPU index {GPU_IDX} not available."
-#     device = torch.device(f"cuda:{GPU_IDX}")
-#     print(f"Using GPU: {GPU_IDX}")
-# else:
-#     print("CUDA is not available. Using CPU instead.")
-#     device = torch.device("cpu")
+nltk.download("punkt")
+# autodetect the available device
+GPU_IDX = 1  # which GPU to use
+if torch.cuda.is_available():
+    num_gpus = torch.cuda.device_count()
+    print(f"Number of available GPUs: {num_gpus}")
+    assert GPU_IDX < num_gpus, f"GPU index {GPU_IDX} not available."
+    device = torch.device(f"cuda:{GPU_IDX}")
+    print(f"Using GPU: {GPU_IDX}")
+else:
+    print("CUDA is not available. Using CPU instead.")
+    device = torch.device("cpu")
 
-# # Configuration for models and their adapters
-# model_config = {
-#     "Base Model": "polygraf-ai/poly-humanizer-base",
-#     "Large Model": "polygraf-ai/poly-humanizer-large",
-#     # "XL Model": {
-#     #     "path": "google/flan-t5-xl",
-#     #     "adapters": {
-#     #         "XL Model Adapter": "polygraf-ai/poly-humanizer-XL-adapter",
-#     #         "XL Law Model Adapter": "polygraf-ai/poly-humanizer-XL-law-adapter",
-#     #         "XL Marketing Model Adapter": "polygraf-ai/marketing-cleaned-13K-grad-acum-4-full",
-#     #         "XL Child Style Model Adapter": "polygraf-ai/poly-humanizer-XL-children-adapter-checkpoint-4000",
-#     #     },
-#     # },
-# }
+# Configuration for models and their adapters
+model_config = {
+    "Base Model": "polygraf-ai/poly-humanizer-base",
+    "Large Model": "polygraf-ai/poly-humanizer-large",
+    "XL Model": {
+        "path": "google/flan-t5-xl",
+        "adapters": {
+            "XL Model Adapter": "polygraf-ai/poly-humanizer-XL-adapter",
+            "XL Law Model Adapter": "polygraf-ai/poly-humanizer-XL-law-adapter",
+            "XL Marketing Model Adapter": "polygraf-ai/marketing-cleaned-13K-grad-acum-4-full",
+            "XL Child Style Model Adapter": "polygraf-ai/poly-humanizer-XL-children-adapter-checkpoint-4000",
+        },
+    },
+}
 
-# # cache the base models, tokenizers, and adapters
-# models, tokenizers = {}, {}
-# for name, config in model_config.items():
-#     path = config if isinstance(config, str) else config["path"]
-#     # initialize model and tokenizer
-#     model = T5ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.bfloat16).to(device)
-#     models[name] = model
-#     tokenizers[name] = T5Tokenizer.from_pretrained(path)
-#     # load all avalable adapters, each being additional roughly 150M parameters
-#     if isinstance(config, dict) and "adapters" in config:
-#         for adapter_name, adapter_path in config["adapters"].items():
-#             model.load_adapter(adapter_path, adapter_name=adapter_name)
-#             print(f"Loaded adapter: {adapter_name}, Num. params: {model.num_parameters()}")
+# cache the base models, tokenizers, and adapters
+models, tokenizers = {}, {}
+for name, config in model_config.items():
+    path = config if isinstance(config, str) else config["path"]
+    # initialize model and tokenizer
+    model = T5ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.bfloat16).to(device)
+    models[name] = model
+    tokenizers[name] = T5Tokenizer.from_pretrained(path)
+    # load all avalable adapters, each being additional roughly 150M parameters
+    if isinstance(config, dict) and "adapters" in config:
+        for adapter_name, adapter_path in config["adapters"].items():
+            model.load_adapter(adapter_path, adapter_name=adapter_name)
+            print(f"Loaded adapter: {adapter_name}, Num. params: {model.num_parameters()}")
 
 
-# def paraphrase_text(
-#     text,
-#     model_name="Base Model",
-#     temperature=1.2,
-#     repetition_penalty=1.0,
-#     top_k=50,
-#     length_penalty=1.0,
-# ):
-#     # select the model, tokenizer and adapter
-#     if "XL" in model_name:  # dynamic adapter load/unload for XL models
-#         # all adapter models use the XL model as the base
-#         tokenizer, model = tokenizers["XL Model"], models["XL Model"]
-#         # set the adapter if it's not already set
-#         if model.active_adapters() != [f"{model_name} Adapter"]:
-#             model.set_adapter(f"{model_name} Adapter")
-#             print(f"Using adapter: {model_name} Adapter")
-#     else:
-#         tokenizer = tokenizers[model_name]
-#         model = models[model_name]
+def paraphrase_text(
+    text,
+    model_name="Base Model",
+    temperature=1.2,
+    repetition_penalty=1.0,
+    top_k=50,
+    length_penalty=1.0,
+):
+    # select the model, tokenizer and adapter
+    if "XL" in model_name:  # dynamic adapter load/unload for XL models
+        # all adapter models use the XL model as the base
+        tokenizer, model = tokenizers["XL Model"], models["XL Model"]
+        # set the adapter if it's not already set
+        if model.active_adapters() != [f"{model_name} Adapter"]:
+            model.set_adapter(f"{model_name} Adapter")
+            print(f"Using adapter: {model_name} Adapter")
+    else:
+        tokenizer = tokenizers[model_name]
+        model = models[model_name]
 
-#     # paraphrase each chunk of text
-#     sentences = sent_tokenize(text)  # sentence boundary detection
-#     paraphrases = []
-#     for sentence in tqdm(sentences):
-#         sentence = sentence.strip()
-#         if len(sentence) == 0:
-#             continue
-#         inputs = tokenizer("Please paraphrase this sentence: " + sentence, return_tensors="pt").to(device)
-#         outputs = model.generate(
-#             **inputs,
-#             do_sample=True,
-#             temperature=temperature,
-#             repetition_penalty=repetition_penalty,
-#             max_length=128,
-#             top_k=top_k,
-#             length_penalty=length_penalty,
-#         )
-#         paraphrased_sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
-#         paraphrases.append(paraphrased_sentence)
-#         print(f"\nOriginal: {sentence}")
-#         print(f"Paraphrased: {paraphrased_sentence}")
+    # paraphrase each chunk of text
+    sentences = sent_tokenize(text)  # sentence boundary detection
+    paraphrases = []
+    for sentence in tqdm(sentences):
+        sentence = sentence.strip()
+        if len(sentence) == 0:
+            continue
+        inputs = tokenizer("Please paraphrase this sentence: " + sentence, return_tensors="pt").to(device)
+        outputs = model.generate(
+            **inputs,
+            do_sample=True,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            max_length=128,
+            top_k=top_k,
+            length_penalty=length_penalty,
+        )
+        paraphrased_sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        paraphrases.append(paraphrased_sentence)
+        print(f"\nOriginal: {sentence}")
+        print(f"Paraphrased: {paraphrased_sentence}")
 
-#     combined_paraphrase = " ".join(paraphrases)
-#     return combined_paraphrase
+    combined_paraphrase = " ".join(paraphrases)
+    return combined_paraphrase