Update app.py

app.py

@@ -1,15 +1,22 @@
 import gradio as gr
 import torch
-from transformers import pipeline, GPTJForCausalLM
+from transformers import pipeline, GPTJForCausalLM, PeftModel, PeftConfig, AutoModelForCausalLM
 
-# load fp 16 model
-model = GPTJForCausalLM.from_pretrained("models/hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
 
-config = AutoConfig.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es", name_or_path="adapter_model.bin")
+config = PeftConfig.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
+model = AutoModelForCausalLM.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es", return_dict=True, load_in_8bit=True, device_map='auto')
+
+# # load fp 16 model
+# model = GPTJForCausalLM.from_pretrained("models/hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
+
+# config = AutoConfig.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es", name_or_path="adapter_model.bin")
 
 # load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
 
+# Load the Lora model
+model = PeftModel.from_pretrained(model, "hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
+
 # create pipeline
 pipe = pipeline("text-generation", model=model, config=config, tokenizer=tokenizer, device=0,)