Update app.py

app.py

@@ -3,6 +3,7 @@ import peft
 from peft import LoraConfig
 from transformers import AutoTokenizer,BitsAndBytesConfig, AutoModelForCausalLM, CLIPVisionModel, AutoProcessor
 import torch
+from peft import PeftModel
 
 clip_model_name = "openai/clip-vit-base-patch32"
 phi_model_name  = "microsoft/phi-2"
@@ -17,39 +18,11 @@ phi_embed  = 2560
 # models
 clip_model = CLIPVisionModel.from_pretrained(clip_model_name).to(device)
 projection = torch.nn.Linear(clip_embed, phi_embed).to(device)
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,)
-
-phi_model = AutoModelForCausalLM.from_pretrained(
-    phi_model_name,
-    torch_dtype=torch.float32,
-    quantization_config=bnb_config,
-    trust_remote_code=True
-)
-lora_alpha = 16
-lora_dropout = 0.1
-lora_r = 64
-peft_config = LoraConfig(
-    lora_alpha=lora_alpha,
-    lora_dropout=lora_dropout,
-    r=lora_r,
-    bias="none",
-    task_type="CAUSAL_LM",
-    target_modules=[
-        "q_proj",
-        'k_proj',
-        'v_proj',
-        'fc1',
-        'fc2'
-    ]
-)
-peft_model = peft.get_peft_model(phi_model, peft_config).to(device)
+phi_model = AutoModelForCausalLM.from_pretrained(phi_model_name,trust_remote_code=True).to(device)
 
 # load weights
-model_to_merge = peft_model.from_pretrained(phi_model,'./model_chkpt/lora_adaptor')
-merged_model = model_to_merge.merge_and_unload()
+model_to_merge = PeftModel.from_pretrained(phi_model,'./model_chkpt/lora_adaptor')
+merged_model   = model_to_merge.merge_and_unload()
 projection.load_state_dict(torch.load('./model_chkpt/step2_projection.pth'))
 
 def model_generate_ans(img,val_q):