Update app.py

app.py

@@ -14,10 +14,8 @@ model_names = {
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the default model initially
-current_model_name = "DialoGPT-medium"
-model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
-model.load_state_dict(torch.load(model_names[current_model_name], map_location=device))
-model.to(device)
+current_model_name = "DialoGPT-med-FT"
+model = AutoModelForCausalLM.from_pretrained(model_names[current_model_name]).to(device)
 
 def load_model(model_name):
     global model, current_model_name
@@ -69,7 +67,7 @@ demo = gr.ChatInterface(
     respond,
     type='messages',
     additional_inputs=[
-        gr.Dropdown(choices=["DialoGPT-med-FT", "DialoGPT-medium"], value="DialoGPT-medium", label="Model"),
+        gr.Dropdown(choices=["DialoGPT-med-FT", "DialoGPT-medium"], value="DialoGPT-med-FT", label="Model"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),