Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
 import airllm
 import os
-from airllm import AirLLMLlama2 
+from airllm.auto_model import AutoModel
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -31,7 +31,7 @@ model = AutoModelForCausalLM.from_pretrained(
 #model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Initialize AirLLM
-air_llm = airllm.AirLLMLlama2(model, tokenizer)
+air_llm = airllm.AutoModel(model, tokenizer)
 
 # Streamlit app configuration
 st.set_page_config(