Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 import streamlit as st
 import os
 from dotenv import load_dotenv
-from airllm import AirLLMInternLM
+from airllm import AutoModel
 from transformers import AutoTokenizer, GenerationConfig
 
 # Load environment variables
@@ -14,7 +14,7 @@ api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 # Initialize model and tokenizer
 MAX_LENGTH = 128
 model_name = "internlm/internlm2_5-7b"
-model = AirLLMInternLM(model_name)
+model = AutoModel.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Streamlit app configuration