Update app.py

app.py

@@ -1,8 +1,15 @@
 import gradio as gr
 
-from transformers import pipeline
+from transformers import pipeline, GPTJForCausalLM
 
-pipe = pipeline("text2text-generation", model="hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
+# load fp 16 model
+model = GPTJForCausalLM.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es", revision="float16", torch_dtype=torch.float16)
+
+# load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es")
+
+# create pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 def predict(text):
   return pipe(f"text: {text}, entities:")["generated_text"]