Upload folder using huggingface_hub

app.py

@@ -33,6 +33,10 @@ def main():
 
     st.title("SongLift LyrGen2")
     st.sidebar.markdown(f"**Model:** {Settings.LLM_MODEL}")
+    if st.sidebar.button("New Song"):
+        st.session_state.chat_history = []
+        st.session_state.current_lyrics = None
+        st.rerun()
 
     # Only run startup once per session
     if 'initialized' not in st.session_state:
@@ -91,36 +95,28 @@ def main():
 
                 # Display sources with content
                 with st.expander("View Sources and Context"):
-                    st.write("### Top Retrieved Contexts")
+                    # Show top retrieved contexts with snippets
+                    st.write("### Retrieved Contexts")
                     for detail in response["context_details"]:
                         st.write(
-                            f"\n**{detail['artist']} - {detail['song']}** "
-                            f"(Similarity: {detail['similarity']}%)"
+                            f"\n**{detail['artist']} - {detail['song']}**"
                         )
-                        st.write("Content snippet:")
                         st.text(detail['content'])
                         st.write("---")
 
-                    st.write("\n### All Similar Sources")
+                    # Show all unique source songs from the chain
+                    st.write("### All Sources Used")
                     seen_sources = set()
-                    unique_sources = []
-
-                    for doc, score in response["source_documents_with_scores"]:
+                    source_docs = response.get("source_documents", [])
+                    for doc in source_docs:
                         source_key = (
-                            doc.metadata['artist'],
-                            doc.metadata['song_title']
+                            doc.metadata.get('artist', 'Unknown'),
+                            doc.metadata.get('song_title', 'Unknown')
                         )
                         if source_key not in seen_sources:
                             seen_sources.add(source_key)
-                            unique_sources.append((doc, score))
-
-                    for doc, score in unique_sources:
-                        similarity = round((1 - score) * 100, 2)
-                        st.write(
-                            f"- {doc.metadata['artist']} - "
-                            f"{doc.metadata['song_title']} "
-                            f"(Similarity: {similarity}%)"
-                        )
+                            st.write(f"- {source_key[0]} - {source_key[1]}")
+                    st.write(f"\n*{len(seen_sources)} unique songs from {len({s[0] for s in seen_sources})} artists*")
 
                 # Update chat history
                 st.session_state.chat_history.append((user_input, lyrics))

src/generator/generator.py

@@ -359,29 +359,9 @@ User Request: {question}"""
             print("Starting lyrics generation process...")
             print(f"Using OpenAI model: {Settings.LLM_MODEL}")
 
-            # Get source documents
-            print("Searching for similar documents...")
-            try:
-                results = self._similarity_search_with_retry(prompt)
-                # results is a list of (Document, score) tuples
-                docs_and_scores = [(doc[0], doc[1]) for doc in results]  # Unpack tuples correctly
-            except Exception as e:
-                print(f"Error during similarity search: {str(e)}")
-                raise RuntimeError(f"Failed to search vector store: {str(e)}")
-
-            # Create detailed context log
-            context_details = []
-            for doc, score in docs_and_scores[:5]:  # Log top 5 for brevity
-                context_details.append({
-                    'artist': doc.metadata['artist'],
-                    'song': doc.metadata['song_title'],
-                    'similarity': f"{score:.2f}",
-                    'content': doc.page_content[:200] + "..."
-                })
-
             try:
                 print("Attempting OpenAI API call...")
-                # Generate response using invoke
+                # Generate response using invoke — DiverseRetriever handles retrieval
                 response = self.qa_chain.invoke({
                     "question": prompt,
                     "chat_history": chat_history
@@ -407,8 +387,19 @@ User Request: {question}"""
                 else:
                     raise RuntimeError(f"OpenAI API error: {error_msg}")
 
-            # Add detailed context to response
-            response["source_documents_with_scores"] = docs_and_scores
+            # Build context details from the chain's actual source documents
+            source_docs = response.get("source_documents", [])
+            context_details = []
+            for doc in source_docs[:10]:
+                context_details.append({
+                    'artist': doc.metadata.get('artist', 'Unknown'),
+                    'song': doc.metadata.get('song_title', 'Unknown'),
+                    'content': doc.page_content[:200] + "..."
+                })
+
+            unique_artists = len({d['artist'] for d in context_details})
+            print(f"Sources shown: {len(context_details)} chunks from {unique_artists} artists")
+
             response["context_details"] = context_details
 
             return response