Update app.py

app.py

@@ -1,30 +1,25 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the tokenizer and model from local files
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("./", config="config.json")
-    model = AutoModelForCausalLM.from_pretrained("./")
+    model_dir = "./"  # Ensure all files, including `vocab.txt`, are in this directory
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForCausalLM.from_pretrained(model_dir)
     return tokenizer, model
 
-# Initialize the model and tokenizer
-tokenizer, model = load_model()
-
-# Set up Streamlit page configuration
 st.set_page_config(page_title="Legal AI Chatbot", layout="centered")
-
 st.title("Legal AI Chatbot")
-st.write("This chatbot provides responses based on a legal language model.")
 
-# User input
-user_input = st.text_input("Enter your query:")
+st.write("Interact with a legal AI chatbot powered by transformers.")
+
+# Load tokenizer and model
+tokenizer, model = load_model()
 
+# User input for chatbot
+user_input = st.text_input("Enter your legal query:")
 if user_input:
-    # Tokenize and generate response
-    inputs = tokenizer.encode(user_input, return_tensors="pt")
-    outputs = model.generate(inputs, max_length=150, num_return_sequences=1)
-    
-    # Decode and display the output
+    inputs = tokenizer(user_input, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=150)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     st.text_area("Response:", response, height=200)