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 AutoModel
+from airllm import AutoModel,AirLLMInternLM
 from transformers import AutoTokenizer, GenerationConfig
 
 # Load environment variables
@@ -14,7 +14,13 @@ api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 # Initialize model and tokenizer
 MAX_LENGTH = 128
 model_name = "internlm/internlm2_5-7b"
-model = AutoModel.from_pretrained(model_name)
+try:
+    model = AirLLMInternLM.from_pretrained(model_name)
+except:
+    try:
+        model=AirLLMInternLM(model_name)
+    except:
+        model=AutoModel.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Streamlit app configuration