Update app.py

app.py

@@ -1,5 +1,5 @@
 ## app.py ##
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from gradio import Interface
 import gradio as gr
 
@@ -8,14 +8,32 @@ MODELS = {
     "T5": "lmsys/fastchat-t5-3b-v1.0",
     "Bert": "bert-base-multilingual-cased",
     "GPT2": "datificate/gpt2-small-spanish",
-    "bloom":"bigscience/bloom"
 }
 
+TOKENIZERS = {
+    "T5": None,
+    "Bert": None,
+    "GPT2": None,
+}
+
+# Load Bloom model separately with memory optimizations
+model_bloom = AutoModelForCausalLM.from_pretrained("bigscience/bloom", low_cpu_mem_usage=True)
+tokenizer_bloom = AutoTokenizer.from_pretrained("bigscience/bloom")
+
 # Define your function
 def generate_and_analyze(model_name, input_text):
     # Load the model from the dictionary using the selected model name
-    model = MODELS[model_name]
-    text_generator = pipeline('text-generation', model=model, device=0) # Use GPU if available
+    if model_name == "bloom":
+        model = model_bloom
+        tokenizer = tokenizer_bloom
+    else:
+        model = MODELS[model_name]
+        tokenizer = TOKENIZERS[model_name]
+        if tokenizer is None:  # Load tokenizer if not already done
+            tokenizer = AutoTokenizer.from_pretrained(model)
+            TOKENIZERS[model_name] = tokenizer
+
+    text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0) # Use GPU if available
     result = text_generator(input_text, max_length=250, do_sample=True)[0]
     return result['generated_text']
 
@@ -23,10 +41,9 @@ def generate_and_analyze(model_name, input_text):
 iface = gr.Interface(
     fn=generate_and_analyze, 
     inputs=[
-        gr.inputs.Dropdown(choices=list(MODELS.keys()), label="Model"),
+        gr.inputs.Dropdown(choices=list(MODELS.keys()) + ["bloom"], label="Model"),
         gr.inputs.Textbox(lines=2, label="Input Text")
     ], 
     outputs="text"
 )
 iface.launch()
-