req file

app.py

@@ -1,81 +1,121 @@
-import gradio as gr
-from sentence_transformers import SentenceTransformer
+import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
 import faiss
-from transformers import pipeline
 import numpy as np
-import os
 
-# File paths
-INDEX_FILE = 'ammons_muse_index.faiss'
-EMBEDDINGS_FILE = 'ammons_muse_embeddings.npy'
-CHUNKS_FILE = 'ammons_muse_chunks.npy'
-TEXT_FILE = 'ammons_muse.txt'
+@st.cache_resource
+def load_models():
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+        embedding_model = AutoModel.from_pretrained("distilbert-base-uncased")
+        generation_model = AutoModelForCausalLM.from_pretrained("gpt2")
+        return tokenizer, embedding_model, generation_model
+    except Exception as e:
+        st.error(f"Error loading models: {str(e)}")
+        return None, None, None
 
-# Load and prepare the text
-def prepare_text():
-    with open(TEXT_FILE, 'r', encoding='utf-8') as file:
-        text = file.read()
-    chunk_size = 1000
-    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
+@st.cache_data
+def load_and_process_text(file_path):
+    try:
+        with open(file_path, 'r', encoding='utf-8') as file:
+            text = file.read()
+        chunks = [text[i:i+512] for i in range(0, len(text), 512)]
+        return chunks
+    except Exception as e:
+        st.error(f"Error loading text file: {str(e)}")
+        return []
 
-# Create or load embeddings and index
-def get_embeddings_and_index(chunks):
-    if os.path.exists(INDEX_FILE) and os.path.exists(EMBEDDINGS_FILE):
-        print("Loading existing index and embeddings...")
-        index = faiss.read_index(INDEX_FILE)
-        embeddings = np.load(EMBEDDINGS_FILE)
-    else:
-        print("Creating new index and embeddings...")
-        model = SentenceTransformer('all-MiniLM-L6-v2')
-        embeddings = model.encode(chunks)
-        dimension = embeddings.shape[1]
-        index = faiss.IndexFlatL2(dimension)
-        index.add(embeddings.astype('float32'))
-        
-        # Save index and embeddings
-        faiss.write_index(index, INDEX_FILE)
-        np.save(EMBEDDINGS_FILE, embeddings)
+@st.cache_data
+def create_embeddings(chunks, tokenizer, embedding_model):
+    embeddings = []
+    for chunk in chunks:
+        inputs = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        with torch.no_grad():
+            outputs = embedding_model(**inputs)
+        embeddings.append(outputs.last_hidden_state.mean(dim=1).squeeze().numpy())
+    return np.array(embeddings)
+
+@st.cache_resource
+def create_faiss_index(embeddings):
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+    return index
+
+def generate_response(query, tokenizer, generation_model, embedding_model, index, chunks):
+    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = embedding_model(**inputs)
+    query_embedding = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
+    
+    k = 3
+    _, I = index.search(query_embedding.reshape(1, -1), k)
+    
+    context = " ".join([chunks[i] for i in I[0]])
+    
+    prompt = f"As the Muse of A.R. Ammons, respond to this query: {query}\nContext: {context}\nMuse:"
     
-    return embeddings, index
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    output = generation_model.generate(input_ids, max_length=200, 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
+
+# Streamlit UI
+st.set_page_config(page_title="A.R. Ammons' Muse Chatbot", page_icon="🎭")
 
-# Load or create chunks
-if os.path.exists(CHUNKS_FILE):
-    chunks = np.load(CHUNKS_FILE, allow_pickle=True).tolist()
-else:
-    chunks = prepare_text()
-    np.save(CHUNKS_FILE, np.array(chunks, dtype=object))
+st.title("A.R. Ammons' Muse Chatbot 🎭")
+st.markdown("""
+    <style>
+    .big-font {
+        font-size:20px !important;
+        font-weight: bold;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+st.markdown('<p class="big-font">Chat with the Muse of A.R. Ammons. Ask questions or discuss poetry!</p>', unsafe_allow_html=True)
 
-# Get embeddings and index
-embeddings, index = get_embeddings_and_index(chunks)
+# 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, tokenizer, embedding_model)
+    index = create_faiss_index(embeddings)
 
-# Set up text generation pipeline
-generator = pipeline('text-generation', model='gpt2')
+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.")
+    st.stop()
 
-# Retrieval function
-def retrieve_relevant_chunks(query, top_k=3):
-    model = SentenceTransformer('all-MiniLM-L6-v2')
-    query_vector = model.encode([query])
-    _, indices = index.search(query_vector.astype('float32'), top_k)
-    return [chunks[i] for i in indices[0]]
+# Initialize chat history
+if 'messages' not in st.session_state:
+    st.session_state.messages = []
 
-# Character response generation
-def generate_character_response(query):
-    relevant_chunks = retrieve_relevant_chunks(query)
-    prompt = f"""As the Muse from A.R. Ammons' poetry, respond to this query:
-Context: {' '.join(relevant_chunks)}
-User: {query}
-Muse:"""
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-    response = generator(prompt, max_length=150, num_return_sequences=1)[0]['generated_text']
-    return response.split('Muse:')[-1].strip()
+# React to user input
+if prompt := st.chat_input("What would you like to ask the Muse?"):
+    st.chat_message("user").markdown(prompt)
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    with st.spinner("The Muse is contemplating..."):
+        try:
+            response = generate_response(prompt, tokenizer, generation_model, embedding_model, index, chunks)
+        except Exception as e:
+            response = f"I apologize, but I encountered an error: {str(e)}"
+    
+    with st.chat_message("assistant"):
+        st.markdown(response)
+    st.session_state.messages.append({"role": "assistant", "content": response})
 
-# Gradio interface
-iface = gr.Interface(
-    fn=generate_character_response,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
-    outputs="text",
-    title="A.R. Ammons' Muse Chatbot",
-    description="Ask a question and get a response from the Muse of A.R. Ammons' poetry."
-)
+# Add a button to clear chat history
+if st.button("Clear Chat History"):
+    st.session_state.messages = []
+    st.experimental_rerun()
 
-iface.launch()
+# Add a footer
+st.markdown("---")
+st.markdown("*Powered by the spirit of A.R. Ammons and the magic of AI*")

requirements.txt

@@ -0,0 +1,8 @@
+
+torch
+transformers
+sentence-transformers
+faiss-cpu
+numpy
+datasets
+streamlit