Metrics and inference in progress

santiatpml · santiatpml · commit 6c47dbef4bc9 · 2023-03-14T17:52:30.000Z
diff --git a/pgml-extension/src/llm_fine_tuning/.gitignore b/pgml-extension/src/llm_fine_tuning/.gitignore
@@ -3,4 +3,7 @@ results
 __pycache__/
 metrics.ipynb
 netflix_titles_medium.csv
-netflix_titles_400.csv
+netflix_titles_400.csv
+*sample.csv
+*.db
+*.ipynb
diff --git a/pgml-extension/src/llm_fine_tuning/generate.py b/pgml-extension/src/llm_fine_tuning/generate.py
@@ -29,7 +29,10 @@
     show_default=True,
 )
 @click.option(
-    "--max_length", default=50, help="Max length of the response", show_default=True
+    "--min_length", default=50, help="Min length of the response", show_default=True
+)
+@click.option(
+    "--max_length", default=50, help="Min length of the response", show_default=True
 )
 @click.option(
     "--num_return_sequences",
@@ -44,21 +47,33 @@
     show_default=True,
 )
 def generate(
-    prompt, model_name, tokenizer_name, max_length, num_return_sequences, temperature
+    prompt, model_name, tokenizer_name, min_length, max_length, num_return_sequences, temperature
 ):
+    cuda_available = torch.cuda.is_available()
     model = AutoModelForCausalLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    if cuda_available:
+        device = "cuda:0"
+    else:
+        device = "cpu"
     generator = pipeline(
-        "text-generation", model=model, tokenizer=tokenizer, max_length=max_length
+        "text-generation", model=model, tokenizer=tokenizer, device=device, max_length = max_length
     )
+    
+    min_length = min(min_length,max_length)
+
     log.info("Prompt: %s" % prompt)
-    log.info(
-        "Generated: %s"
-        % generator(
-            prompt, num_return_sequences=num_return_sequences, temperature=temperature
-        )
+    outputs = generator(
+        prompt,
+        do_sample=True,
+        min_length=min_length,
+        num_return_sequences=num_return_sequences,
+        temperature=temperature,
     )
 
+    for _id, output in enumerate(outputs):
+        log.info("Generated %d: %s" % (_id, output["generated_text"]))
+
 
 if __name__ == "__main__":
     generate()
diff --git a/pgml-extension/src/llm_fine_tuning/metrics.py b/pgml-extension/src/llm_fine_tuning/metrics.py
@@ -44,7 +44,13 @@
     help="Stride length for computing perplexity",
     show_default=True,
 )
-def metrics(filename, column_name, model_name, tokenizer_name, stride):
+@click.option(
+    "--max_length_key",
+    default="n_positions",
+    help="Key in model configuration that maps to max length of the embeddings",
+    show_default=True,
+)
+def metrics(filename, column_name, model_name, tokenizer_name, stride, max_length_key):
     if os.path.exists(filename):
         test = load_dataset("csv", data_files=filename)
     else:
@@ -55,16 +61,25 @@ def metrics(filename, column_name, model_name, tokenizer_name, stride):
 
     device = "cpu"
     if cuda_available:
-        device = "cuda"
-        model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
-    else:
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+        device = "cuda:0"
+    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
 
-    encodings = tokenizer("\n\n".join(test["train"][column_name]), return_tensors="pt")
+    full_text = ""
+    for entry in test["train"][column_name]:
+        if entry:
+            full_text += "\n\n" + entry
+
+    encodings = tokenizer(full_text, return_tensors="pt")
+
+    config = model.config.to_dict()
+    if max_length_key in config.keys():
+        max_length = config[max_length_key]
+    else:
+        log.info("Configuration keys " + ",".join(config.keys()))
+        raise ValueError("%s does not exist in model configuration"%max_length_key)
 
-    max_length = model.config.n_positions
     stride = min(stride, max_length)
     seq_len = encodings.input_ids.size(1)
 
@@ -93,7 +108,7 @@ def metrics(filename, column_name, model_name, tokenizer_name, stride):
             break
 
     ppl = torch.exp(torch.stack(nlls).sum() / end_loc)
-    log.info("Perplexity = %0.3f (lower is better)" % ppl)
+    log.info("Number of parameters = %d, Perplexity = %0.3f (lower is better)" % (model.num_parameters(), ppl))
 
 
 if __name__ == "__main__":
