Update app.py

app.py

@@ -6,12 +6,20 @@ import torch
 import gradio as gr
 
 
-# Load default model as GPT2
+# Load default model as GPT2 and other models
 
 
 tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
 model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
 
+tokenizer_gpt2 = AutoTokenizer.from_pretrained("openai-community/gpt2")
+model_gpt2 = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+
+tokenizer_gemma = AutoTokenizer.from_pretrained("google/gemma-2b")
+model_gemma = AutoModelForCausalLM.from_pretrained("google/gemma-2b")
+
+tokenizer_Mistral = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
+model_Mistral = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
 
 # Define functions
 
@@ -129,13 +137,17 @@ def generate(input_text, number_steps, number_beams, number_beam_groups, diversi
 def load_model(model_selected):
     
     if model_selected == "gpt2":
-        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
-        model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2", pad_token_id=tokenizer.eos_token_id)
+        tokenizer = tokenizer_gpt2
+        model = model_gpt2
         #print (model_selected + " loaded")
     
     if model_selected == "Gemma 2":
-        tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
-        model = AutoModelForCausalLM.from_pretrained("google/gemma-2b")
+        tokenizer = tokenizer_gemma
+        model = model_gemma
+    
+    if model_selected == "Mistral":
+        tokenizer = tokenizer_Mistral
+        model = model_Mistral
 
 
 
@@ -298,7 +310,7 @@ with gr.Blocks() as demo:
         
         with gr.Column (scale=0, min_width=200) as Models_Strategy:
             
-            model_selected = gr.Radio (["gpt2", "Gemma 2"], label="ML Model", value="gpt2")
+            model_selected = gr.Radio (["gpt2", "Gemma 2", "Mistral"], label="ML Model", value="gpt2")
             strategy_selected = gr.Radio (["Sampling", "Beam Search", "Diversity Beam Search","Contrastive"], label="Search strategy", value = "Sampling", interactive=True)