bug fixed

app.py

@@ -1,17 +1,22 @@
 import streamlit as st
+from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 import os.path 
 import pickle
 import torch
 
+base_model_id = "meta-llama/Llama-3.2-3B-Instruct"
 model_id = "HiGenius/Headline-Generation-Model"
+hf_token = os.environ.get('HF_TOKEN')
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 @st.cache_resource
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    base_model = AutoModelForCausalLM.from_pretrained(base_model_id, use_auth_token=hf_token)
+    model = PeftModel.from_pretrained(base_model, model_id).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_auth_token=hf_token)
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side='left'