Update app.py

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
+from mixtrl import MixtralConfig, MixtralForCausalLM
 from peft import PeftModel
 
 # Load the tokenizer and model
 model_name = "JuliaUpton/Math_AI"
 model_architecture = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 mixtral_tokenizer = AutoTokenizer.from_pretrained(model_architecture)
-merged_model = PeftModel.from_pretrained(model_name, model=AutoModelForCausalLM.from_pretrained(model_architecture, device_map="auto"))
+mixtral_config = MixtralConfig.from_pretrained(model_architecture)
+merged_model = PeftModel.from_pretrained(model_name, model=MixtralForCausalLM.from_pretrained(model_architecture, config=mixtral_config, device_map="auto"))
 
 # Define the input formatting function
 def input_from_text(instruction):