Update app.py

app.py

@@ -1,12 +1,20 @@
 import spaces
 import gradio as gr
 import torch
-from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer, AutoConfig
 from peft import PeftModel
 
 MODEL_ID = "unsloth/Meta-Llama-3.1-70B-bnb-4bit"
 ADAPTER_ID = "marcelbinz/Llama-3.1-Centaur-70B-adapter"
 
+cfg = AutoConfig.from_pretrained(MODEL_ID)
+cfg.rope_scaling = {
+    "type": "yarn",
+    "factor": 4.0,
+    "original_max_position_embeddings": 8192,
+}
+cfg.max_position_embeddings = 32768
+
 bnb_4bit_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -19,6 +27,7 @@ model_base = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
     attn_implementation="flash_attention_2",
+    config=cfg,
     quantization_config=bnb_4bit_config,
 )