Update app.py

app.py

@@ -4,6 +4,11 @@ import keras_nlp
 
 import numpy as np
 import pandas as pd
+
+import os
+os.environ["KERAS_BACKEND"] = "jax" # you can also use tensorflow or torch
+os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1" # avoid memory fragmentation on JAX backend.
+
 keras.utils.set_random_seed(42)
 
 gemma_lm = keras_nlp.models.CausalLM.from_preset("hf://soufyane/gemma_2b_instruct_FT_DATA_SCIENCE_lora36_1")