change model to sharded version

app.py

@@ -6,13 +6,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 peft_model_id = "hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es"
 config = PeftConfig.from_pretrained(peft_model_id)
 model = AutoModelForCausalLM.from_pretrained(
-    config.base_model_name_or_path,
+    "DavidFM43/bertin-gpt-j-6b-half-sharded",
     return_dict=True,
     load_in_8bit=True,
     device_map="auto",
-    revision="half",
 )
-model.tie_weights()
 tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
 # load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
@@ -28,9 +26,7 @@ def gen_entities(text):
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=256, eos_token_id=50258)
 
-    response = tokenizer.batch_decode(
-        output_tokens.detach().cpu().numpy(), skip_special_tokens=False
-    )[0]
+    response = tokenizer.batch_decode(output_tokens.detach().cpu().numpy(), skip_special_tokens=False)[0]
 
     return response[response.find("entities") : response.find("<EP>")]