Update app.py

app.py

@@ -20,8 +20,8 @@ model_gpt2 = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
 tokenizer_gemma = AutoTokenizer.from_pretrained("google/gemma-2b")
 model_gemma = AutoModelForCausalLM.from_pretrained("google/gemma-2b", token=token)
 
-tokenizer_Mistral = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
-model_Mistral = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
+tokenizer_gpt-neo = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
+model_gpt-neo = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
 
 # Define functions
 
@@ -138,7 +138,7 @@ def generate(input_text, number_steps, number_beams, number_beam_groups, diversi
 
 def load_model(model_selected):
     
-    if model_selected == "gpt2":
+    if model_selected == "GPT2":
         tokenizer = tokenizer_gpt2
         model = model_gpt2
         #print (model_selected + " loaded")
@@ -147,9 +147,9 @@ def load_model(model_selected):
         tokenizer = tokenizer_gemma
         model = model_gemma
     
-    if model_selected == "Mistral":
-        tokenizer = tokenizer_Mistral
-        model = model_Mistral
+    if model_selected == "Eleuther GPT Neo 1.3B":
+        tokenizer = tokenizer_gpt-neo
+        model = model_gpt-neo
 
 
 
@@ -311,7 +311,7 @@ with gr.Blocks() as demo:
         
         with gr.Column (scale=0, min_width=200) as Models_Strategy:
             
-            model_selected = gr.Radio (["gpt2", "Gemma 2", "Mistral"], label="ML Model", value="gpt2")
+            model_selected = gr.Radio (["GPT2", "Gemma 2", "Eleuther GPT Neo 1.3B"], label="ML Model", value="gpt2")
             strategy_selected = gr.Radio (["Sampling", "Beam Search", "Diversity Beam Search","Contrastive"], label="Search strategy", value = "Sampling", interactive=True)