From 558bcbc93a0a434efd6ae2a77fb811984f150a29 Mon Sep 17 00:00:00 2001
From: Lorenzo Toniazzi <lorenzo.toniazzi@stepstone.com>
Date: Mon, 12 Aug 2024 10:41:24 +0100
Subject: [PATCH] Add assertions

---
 .gitignore                                  |  3 ++
 tests/test_lora_conversion_and_inference.sh | 48 ++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index c9b4d99839729c..754c6737e36c37 100644
--- a/.gitignore
+++ b/.gitignore
@@ -130,3 +130,6 @@ poetry.toml
 
 # Scripts
 !/scripts/install-oneapi.bat
+
+# Test models for lora adapters
+/reduce-llms-for-testing
diff --git a/tests/test_lora_conversion_and_inference.sh b/tests/test_lora_conversion_and_inference.sh
index 5780dcf6858934..60a905b9b1a928 100755
--- a/tests/test_lora_conversion_and_inference.sh
+++ b/tests/test_lora_conversion_and_inference.sh
@@ -12,9 +12,28 @@ else
     echo "Repository already exists. Skipping clone."
 fi
 
+# Load the expected starting strings from the text file
+EXPECTED_BASE=$(cat $MODELS_REPO/data/pale_blue_dot.txt)
+EXPECTED_LORA_HOT=$(cat $MODELS_REPO/data/bohemian_rhapsody.txt)
+EXPECTED_LORA_MERGED=$(cat $MODELS_REPO/data/bohemian_rhapsody.txt)
+
 # Declare a regular array to store results
 results=()
 
+trim_leading_whitespace() {
+    local input_string="$1"
+    # Remove leading whitespace
+    echo "${input_string#"${input_string%%[![:space:]]*}"}"
+}
+
+extract_starting_substring() {
+    local reference_string="$1"
+    local target_string="$2"
+    
+    local target_length=${#target_string}
+    echo "${reference_string:0:$target_length}"
+}
+
 run_conversion_and_inference_lora() {
     local model_name=$1
     local size_matrix=$2
@@ -50,6 +69,33 @@ run_conversion_and_inference_lora() {
     OUTPUT_LORA_MERGED=$(llama-cli -m $MODELS_REPO/$model_name/size=$size_matrix/base/Base-F32-lora-merged.gguf \
     -p "I see a little silhouetto" -n 50 --seed 42 --temp 0)
 
+    # Extract the corresponding substring from EXPECTED_BASE
+    # and remove initial white spaces in OUTPUT_BASE
+    OUTPUT_BASE=$(trim_leading_whitespace "$OUTPUT_BASE")
+    EXPECTED_BASE=$(extract_starting_substring "$EXPECTED_BASE" "$OUTPUT_BASE")
+    OUTPUT_LORA_HOT=$(trim_leading_whitespace "$OUTPUT_LORA_HOT")
+    EXPECTED_LORA_HOT=$(extract_starting_substring "$EXPECTED_LORA_HOT" "$OUTPUT_LORA_HOT")
+    OUTPUT_LORA_MERGED=$(trim_leading_whitespace "$OUTPUT_LORA_MERGED")
+    EXPECTED_LORA_MERGED=$(extract_starting_substring "$EXPECTED_LORA_MERGED" "$OUTPUT_LORA_MERGED")
+
+    # Compare the actual output with the expected start
+    if [[ "$OUTPUT_BASE" != "$EXPECTED_BASE" ]]; then
+        echo "Error: $model_name OUTPUT_BASE does not start with the expected string."
+        echo -e "Out=$OUTPUT_BASE\n\nExp=$EXPECTED_BASE"
+        exit 1
+    fi
+    if [[ "$OUTPUT_LORA_HOT" != "$EXPECTED_LORA_HOT" ]]; then
+        echo "Error: $model_name OUTPUT_LORA_HOT does not start with the expected string."
+        echo -e "Out=$OUTPUT_LORA_HOT\n\nExp=$EXPECTED_LORA_HOT"
+        exit 1
+    fi
+    if [[ "$OUTPUT_LORA_MERGED" != "$EXPECTED_LORA_MERGED" ]]; then
+        echo "Error: $model_name OUTPUT_LORA_MERGED does not start with the expected string."
+        echo -e "Out=$OUTPUT_LORA_MERGED\n\nExp=$EXPECTED_LORA_MERGED"
+        exit 1
+    fi
+
+
     # Store the results in the regular array
     results+=("
     \n\033[1mResults for $model_name with size $size_matrix:\033[0m
@@ -64,7 +110,7 @@ run_conversion_and_inference_lora() {
 
 # Array of parameters to iterate over
 declare -a params=(
-    "Gemma2ForCausalLM 64"
+    # "Gemma2ForCausalLM 64"
     "LlamaForCausalLM 64"
     "Phi3ForCausalLM 64"
 )