Update app.py

app.py

@@ -3,19 +3,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
 import airllm
 import os
-from airllm.auto_model import AutoModel
 from dotenv import load_dotenv
 
 load_dotenv()
 
-
-
-
-
-
-
-api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
 # Retrieve the API token from the environment variables
 api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
@@ -27,11 +18,9 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16,
     use_auth_token=api_token
 )
-# Load the base version of the model
-#model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Initialize AirLLM
-air_llm = airllm.AutoModel(model)
+air_llm = airllm.AutoModel()
 
 # Streamlit app configuration
 st.set_page_config(
@@ -85,7 +74,7 @@ user_input = st.text_input("You: ", "")
 if st.button("Send"):
     if user_input:
         # Generate response using AirLLM
-        response = air_llm.generate_response(user_input)
+        response = air_llm.generate_response(model, tokenizer, user_input)
         st.text_area("Bot:", value=response, height=200, max_chars=None)
     else:
         st.warning("Please enter a message.")