From fa09dc37b5a572ffac01fb761d9be9d2cf7c64eb Mon Sep 17 00:00:00 2001
From: apun <apun@andrew.cmu.edu>
Date: Mon, 31 Mar 2025 10:39:59 -0400
Subject: [PATCH] Add test_finetuned_llm test

---
 tests/test_infer.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/tests/test_infer.py b/tests/test_infer.py
index 2d49c66..a37787f 100644
--- a/tests/test_infer.py
+++ b/tests/test_infer.py
@@ -1,4 +1,4 @@
-from legogpt.models.legogpt import LegoGPT
+from legogpt.models.legogpt import LegoGPT, create_instruction
 from legogpt.models.llm import LLM
 
 
@@ -21,6 +21,20 @@ def test_llm():
     print(prompt + '|' + output + '|' + output_continuation)
 
 
+def test_finetuned_llm():
+    """
+    Tests running the finetuned LegoGPT model with no other guidance (e.g. rejection sampling).
+    """
+    llm = LLM('/data/apun/finetuned_hf/Llama-3.2-1B-Instruct_finetuned_combined_2')
+    messages = [
+        {'role': 'system', 'content': 'You are a helpful assistant.'},
+        {'role': 'user', 'content': create_instruction('A basic chair with four legs.')},
+    ]
+    prompt = llm.tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors='pt')
+    output = llm(prompt, max_new_tokens=8192)
+    print(output)
+
+
 def test_infer():
     """
     Runs LegoGPT inference on a simple prompt.