Adding med fine tuned model to replace small

app.py

@@ -1,24 +1,30 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # Load the shared tokenizer (using a tokenizer from DialoGPT models)
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
 
-# Define the model names
+# Define the model names, including the locally saved fine-tuned model
 model_names = {
-    "DialoGPT-small": "microsoft/DialoGPT-small",
+    "DialoGPT-med-FT": "DialoGPT-med-FT.bin",
     "DialoGPT-medium": "microsoft/DialoGPT-medium"
 }
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Pre-load the models
 loaded_models = {
-    model_name: AutoModelForCausalLM.from_pretrained(model_path)
-    for model_name, model_path in model_names.items()
+    "DialoGPT-med-FT": AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 }
+loaded_models["DialoGPT-med-FT"].load_state_dict(torch.load(model_names["DialoGPT-med-FT"]))
+loaded_models["DialoGPT-med-FT"].to(device)
+
+loaded_models["DialoGPT-medium"] = AutoModelForCausalLM.from_pretrained(model_names["DialoGPT-medium"]).to(device)
 
 def respond(
     message,
-    history: list[tuple[str, str]],
+    history: list[dict],
     model_choice,
     max_tokens,
     temperature,
@@ -29,12 +35,12 @@ def respond(
 
     # Prepare the input by concatenating the history into a dialogue format
     input_text = ""
-    for user_msg, bot_msg in history:
-        input_text += f"User: {user_msg}\nAssistant: {bot_msg}\n"
+    for message_pair in history:
+        input_text += f"{message_pair['role']}: {message_pair['content']}\n"
     input_text += f"User: {message}\nAssistant:"
 
     # Tokenize the input text using the shared tokenizer
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(model.device)
 
     # Generate the response using the selected DialoGPT model
     output_tokens = model.generate(
@@ -51,9 +57,10 @@ def respond(
 
 # Define the Gradio interface
 demo = gr.ChatInterface(
+    type='messages',
     respond,
     additional_inputs=[
-        gr.Dropdown(choices=["DialoGPT-small", "DialoGPT-medium"], value="DialoGPT-small", 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)"),