Update app.py

app.py

@@ -1,29 +1,23 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import AutoModel, AutoTokenizer
+import torch
+from transformers import AutoTokenizer
 from adapters import AutoAdapterModel
 
-
 model_name = "unsloth/Meta-Llama-3.1-8B-Instruct"
 
+# Load tokenizer with 4-bit quantization
+tokenizer = AutoTokenizer.from_pretrained(model_name, load_in_4bit=True)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, load_in_4bit = True)
-
-# Load the base model with adapters
-model = AutoAdapterModel.from_pretrained(model_name)
+# Load the base model with adapters, ensuring it's loaded in 4-bit
+model = AutoAdapterModel.from_pretrained(model_name, load_in_4bit=True).to("cpu")
 
+# Load the adapter
 model.load_adapter("Braszczynski/Llama-3.2-3B-Instruct-bnb-4bit-460steps")
 
+# Free up unused memory
+torch.cuda.empty_cache()
 
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
+def respond(message, history, system_message, max_tokens, temperature, top_p):
     # Combine system message and chat history
     chat_history = f"{system_message}\n"
     for user_msg, bot_reply in history:
@@ -33,10 +27,10 @@ def respond(
     # Tokenize the input
     inputs = tokenizer(chat_history, return_tensors="pt", truncation=True).to("cuda")
 
-    # Generate response
+    # Generate response with reduced max tokens if necessary
     outputs = model.generate(
         inputs["input_ids"],
-        max_new_tokens=max_tokens,
+        max_new_tokens=max_tokens,  # Consider setting a lower default
         temperature=temperature,
         top_p=top_p,
         pad_token_id=tokenizer.eos_token_id