Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import os, torch
 from datasets import load_dataset
 from huggingface_hub import HfApi, login
-from peft import LoraConfig
+from peft import AutoPeftModelForCausalLM, LoraConfig
+from random import randint
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, pipeline
 from trl import SFTTrainer, setup_chat_format
 
@@ -21,20 +22,46 @@ base_model_id = "codellama/CodeLlama-7b-hf" # "ibm-granite/granite-8b-code-instr
 dataset = "b-mc2/sql-create-context"
 
 def prompt_model(model_id, system_prompt, user_prompt, schema):
-    pipe = pipeline("text-generation", 
-                    model=model_id, 
-                    model_kwargs={"torch_dtype": torch.bfloat16}, 
-                    device_map="auto",
-                    max_new_tokens=1000)
-    messages = [
-      {"role": "system", "content": system_prompt.format(schema=schema)},
-      {"role": "user", "content": user_prompt},
-      {"role": "assistant", "content": ""}
-    ]
-    output = pipe(messages)
-    result = output[0]["generated_text"][-1]["content"]
-    print(result)
-    return result
+#    pipe = pipeline("text-generation", 
+#                    model=model_id, 
+#                    model_kwargs={"torch_dtype": torch.bfloat16}, 
+#                    device_map="auto",
+#                    max_new_tokens=1000)
+#    messages = [
+#      {"role": "system", "content": system_prompt.format(schema=schema)},
+#      {"role": "user", "content": user_prompt},
+#      {"role": "assistant", "content": ""}
+#    ]
+#    output = pipe(messages)
+#    result = output[0]["generated_text"][-1]["content"]
+#    print(result)
+#    return result
+    
+    peft_model_id = "./code-llama-7b-text-to-sql"
+    # peft_model_id = args.output_dir
+     
+    # Load Model with PEFT adapter
+    model = AutoPeftModelForCausalLM.from_pretrained(
+      peft_model_id,
+      device_map="auto",
+      torch_dtype=torch.float16
+    )
+    tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
+    # load into pipeline
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+    ###
+
+    eval_dataset = load_dataset("json", data_files="test_dataset.json", split="train")
+    rand_idx = randint(0, len(eval_dataset))
+     
+    # Test on sample
+    prompt = pipe.tokenizer.apply_chat_template(eval_dataset[rand_idx]["messages"][:2], tokenize=False, add_generation_prompt=True)
+    outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)
+     
+    print(f"Query:\n{eval_dataset[rand_idx]['messages'][1]['content']}")
+    print(f"Original Answer:\n{eval_dataset[rand_idx]['messages'][2]['content']}")
+    print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")
 
 def fine_tune_model(base_model_id, dataset):
     #tokenizer = download_model(base_model_id)