Fix max_length error and implement data persistence

app.py

@@ -3,6 +3,8 @@ import torch
 from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
 import faiss
 import numpy as np
+import os
+import pickle
 
 @st.cache_resource
 def load_models():
@@ -56,12 +58,27 @@ def generate_response(query, tokenizer, generation_model, embedding_model, index
     prompt = f"As the Muse of A.R. Ammons, respond to this query: {query}\nContext: {context}\nMuse:"
     
     input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output = generation_model.generate(input_ids, max_length=200, num_return_sequences=1, temperature=0.7)
+    output = generation_model.generate(input_ids, max_new_tokens=100, num_return_sequences=1, temperature=0.7)
     response = tokenizer.decode(output[0], skip_special_tokens=True)
     
     muse_response = response.split("Muse:")[-1].strip()
     return muse_response
 
+def save_data(chunks, embeddings, index):
+    with open('chunks.pkl', 'wb') as f:
+        pickle.dump(chunks, f)
+    np.save('embeddings.npy', embeddings)
+    faiss.write_index(index, 'faiss_index.bin')
+
+def load_data():
+    if os.path.exists('chunks.pkl') and os.path.exists('embeddings.npy') and os.path.exists('faiss_index.bin'):
+        with open('chunks.pkl', 'rb') as f:
+            chunks = pickle.load(f)
+        embeddings = np.load('embeddings.npy')
+        index = faiss.read_index('faiss_index.bin')
+        return chunks, embeddings, index
+    return None, None, None
+
 # Streamlit UI
 st.set_page_config(page_title="A.R. Ammons' Muse Chatbot", page_icon="🎭")
 
@@ -79,9 +96,12 @@ st.markdown('<p class="big-font">Chat with the Muse of A.R. Ammons. Ask question
 # Load models and data
 with st.spinner("Loading models and data..."):
     tokenizer, embedding_model, generation_model = load_models()
-    chunks = load_and_process_text('ammons_muse.txt')
-    embeddings = create_embeddings(chunks, embedding_model)
-    index = create_faiss_index(embeddings)
+    chunks, embeddings, index = load_data()
+    if chunks is None or embeddings is None or index is None:
+        chunks = load_and_process_text('ammons_muse.txt')
+        embeddings = create_embeddings(chunks, embedding_model)
+        index = create_faiss_index(embeddings)
+        save_data(chunks, embeddings, index)
 
 if tokenizer is None or embedding_model is None or generation_model is None or not chunks:
     st.error("Failed to load necessary components. Please try again later.")