Update app.py

app.py

@@ -1,37 +1,30 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import streamlit as st
-import airllm
 import os
 from dotenv import load_dotenv
+from airllm import AutoModel
 
+# Load environment variables
 load_dotenv()
 
 # Retrieve the API token from the environment variables
 api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-# Load GEMMA 27B model and tokenizer using the API token
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it", use_auth_token=api_token)
-model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-2-9b-it",
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-    use_auth_token=api_token
-)
-
-# Initialize AirLLM
-air_llm = airllm.AutoModel()
+# Initialize model and tokenizer using the AutoModel from AirLLM
+MAX_LENGTH = 128
+model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct")
 
 # Streamlit app configuration
 st.set_page_config(
-    page_title="Chatbot with GEMMA 27B and AirLLM",
+    page_title="Conversational Chatbot with Platypus2-70B and AirLLM",
     page_icon="🤖",
     layout="wide",
     initial_sidebar_state="expanded",
 )
 
 # App title
-st.title("Conversational Chatbot with GEMMA 27B and AirLLM")
+st.title("Conversational Chatbot with Platypus2-70B and AirLLM")
 
 # Sidebar configuration
 st.sidebar.header("Chatbot Configuration")
@@ -73,8 +66,23 @@ elif theme == "Light":
 user_input = st.text_input("You: ", "")
 if st.button("Send"):
     if user_input:
-        # Generate response using AirLLM
-        response = air_llm.generate_response(model, tokenizer, user_input)
+        # Tokenize user input
+        input_tokens = model.tokenizer(user_input,
+            return_tensors="pt", 
+            return_attention_mask=False, 
+            truncation=True, 
+            max_length=MAX_LENGTH, 
+            padding=False)
+        
+        # Generate response
+        generation_output = model.generate(
+            input_tokens['input_ids'].cuda(), 
+            max_new_tokens=20,
+            use_cache=True,
+            return_dict_in_generate=True)
+        
+        # Decode response
+        response = model.tokenizer.decode(generation_output.sequences[0])
         st.text_area("Bot:", value=response, height=200, max_chars=None)
     else:
         st.warning("Please enter a message.")
@@ -83,6 +91,6 @@ if st.button("Send"):
 st.sidebar.markdown(
     """
     ### About
-    This is a conversational chatbot built using the base version of the GEMMA 27B model and AirLLM.
+    This is a conversational chatbot built using the Platypus2-70B model and AirLLM.
     """
 )