Update app.py

app.py

@@ -136,11 +136,14 @@ def generate(input_text, number_steps, number_beams, number_beam_groups, diversi
     
 #--------ON SELECTING MODEL------------------------
 
-def load_model(model_selected):
+def select_model(model_selected):
+    
+    global model_name
     
     if model_selected == "GPT2":
-        tokenizer = tokenizer_gpt2
-        model = model_gpt2
+        model_name = "openai-community/gpt2"
+        #tokenizer = tokenizer_gpt2
+        #model = model_gpt2
         #print (model_selected + " loaded")
     
     #if model_selected == "Gemma 2":
@@ -148,10 +151,19 @@ def load_model(model_selected):
         #model = model_gemma
     
     if model_selected == "Qwen2":
-        tokenizer = tokenizer_qwen
-        model = model_qwen
+        model_name = "Qwen/Qwen2-0.5B"
+        #tokenizer = tokenizer_qwen
+        #model = model_qwen
+
+# On clicking load button
 
+def load_model ():
+    
+    global model_name
+    tokenizer_gpt2 = AutoTokenizer.from_pretrained(model_name)
+    model_gpt2 = AutoModelForCausalLM.from_pretrained(model_name)
 
+    
 
 #--------ON SELECT NO. OF RETURN SEQUENCES----------
 
@@ -313,7 +325,7 @@ with gr.Blocks() as demo:
             
             model_selected = gr.Radio (["GPT2", "Qwen2"], label="ML Model", value="GPT2")
             strategy_selected = gr.Radio (["Sampling", "Beam Search", "Diversity Beam Search","Contrastive"], label="Search strategy", value = "Sampling", interactive=True)
-            
+            load_model_button = gr.Button("Load")
         
         with gr.Column(scale=1):
                 
@@ -390,10 +402,10 @@ with gr.Blocks() as demo:
 
     
 
-#----------ON SELECTING/CHANGING: RETURN SEEQUENCES/NO OF BEAMS/BEAM GROUPS/TEMPERATURE--------
+    #----------ON SELECTING/CHANGING: RETURN SEEQUENCES/NO OF BEAMS/BEAM GROUPS/TEMPERATURE--------
             
     model_selected.change(
-        fn=load_model, inputs=[model_selected], outputs=[]
+        fn=select_model, inputs=[model_selected], outputs=[]
     )
     
     #num_return_sequences.change(
@@ -421,6 +433,12 @@ with gr.Blocks() as demo:
         outputs=[out_markdown]
     )
 
+    load_model_button.click(
+        fn=load_model,
+        inputs=[model_selected],
+        outputs=[]
+    )
+
 
     with gr.Row():