Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import pandas as pd
 from datetime import datetime, timedelta, timezone
 import torch
-from config import hugging_face_token, init_google_sheets_client, models, quantized_models, default_model_name, user_names, google_sheets_name, MAX_INTERACTIONS
+from config import hugging_face_token, init_google_sheets_client, models, default_model_name, user_names, google_sheets_name, MAX_INTERACTIONS
 import spaces
 
 # Hack for ZeroGPU
@@ -48,28 +48,14 @@ def load_model(model_name):
             del model
             torch.cuda.empty_cache()
         
-        tokenizer = AutoTokenizer.from_pretrained(
-            models[model_name], 
-            padding_side='left', 
-            token=hugging_face_token, 
-            trust_remote_code=True
-        )
+        tokenizer = AutoTokenizer.from_pretrained(models[model_name], padding_side='left', token=hugging_face_token, trust_remote_code=True)
         
         # Ensure the padding token is set
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
             tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
         
-        model = AutoModelForCausalLM.from_pretrained(
-            models[model_name], 
-            token=hugging_face_token, 
-            trust_remote_code=True
-        )
-        
-        # Only move to CUDA if it's not a quantized model
-        if model_name not in quantized_models:
-            model = model.to("cuda")
-        
+        model = AutoModelForCausalLM.from_pretrained(models[model_name], token=hugging_face_token, trust_remote_code=True).to("cuda")
         selected_model = model_name
     except Exception as e:
         print(f"Error loading model {model_name}: {e}")
@@ -90,12 +76,6 @@ def interact(user_input, history, interaction_count):
         if tokenizer is None or model is None:
             raise ValueError("Tokenizer or model is not initialized.")
         
-        # Determine the device to use (either CUDA if available, or CPU)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Ensure the model is on the correct device
-        model.to(device)
-        
         if interaction_count >= MAX_INTERACTIONS:
             user_input += ". Thank you for your questions. Our session is now over. Goodbye!"
         
@@ -108,8 +88,8 @@ def interact(user_input, history, interaction_count):
         
         prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         
-        # Move input tensor to the same device as the model
-        input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)
+        # Generate response using selected model
+        input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to("cuda")
         chat_history_ids = model.generate(input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id, temperature=0.1)
         response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
         
@@ -202,9 +182,6 @@ def load_user_guide():
     with open('user_guide.txt', 'r') as file:
         return file.read()
 
-# Combine both model dictionaries
-all_models = {**models, **quantized_models}
-
 # Create the chat interface using Gradio Blocks
 with gr.Blocks() as demo:
     with gr.Tabs():
@@ -213,7 +190,7 @@ with gr.Blocks() as demo:
 
             gr.Markdown("## Context")
             with gr.Group():
-                model_dropdown = gr.Dropdown(choices=list(all_models.keys()), label="Select Model", value=default_model_name)
+                model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model", value=selected_model)
                 user_dropdown = gr.Dropdown(choices=user_names, label="Select User Name")
                 initial_story = stories[0]["title"] if stories else None
                 story_dropdown = gr.Dropdown(choices=[story["title"] for story in stories], label="Select Story", value=initial_story)
@@ -244,4 +221,4 @@ with gr.Blocks() as demo:
     send_message_button.click(fn=interact, inputs=[chatbot_input, chat_history_json, interaction_count], outputs=[chatbot_input, chatbot_output, chat_history_json, interaction_count])
     save_button.click(fn=save_comment_score, inputs=[chatbot_output, score_input, comment_input, story_dropdown, user_dropdown, system_prompt_dropdown], outputs=[data_table, comment_input])
 
-demo.launch()
+demo.launch()