Add test_finetuned_llm test

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.