Update app.py

app.py

@@ -1,74 +1,86 @@
 import gradio as gr
+import chromadb
+from sentence_transformers import SentenceTransformer
 from llama_cpp import Llama
 
-# Load the Llama model
+# ✅ Initialize ChromaDB
+chroma_client = chromadb.PersistentClient(path="./chromadb_store")
+collection = chroma_client.get_or_create_collection(name="curly_strings_knowledge")
+
+# ✅ Load Local Embedding Model
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+
+# ✅ Curly Strings Knowledge (Stored in ChromaDB as Vectors)
+knowledge_base = [
+    {"id": "song_list", "text": """
+    Curly Strings is an Estonian folk band known for blending traditional and modern sounds.
+    Here are some of their popular songs:
+    1. Kalakesed
+    2. Kus mu süda on ...
+    3. Vitsalaul
+    4. Viimases jaamas
+    5. Salaja
+    6. Üle ilma
+    7. Šveits
+    8. Kallimale
+    9. Üksteist peab hoidma
+    10. Suuda öelda ei
+    """},
+    {"id": "related_artists", "text": """
+    If you enjoy Curly Strings, you might also like:
+    - Trad.Attack!
+    - Eesti Raadio laululapsed
+    - Körsikud
+    - Karl-Erik Taukar
+    - Dag
+    """},
+    {"id": "background", "text": """
+    Curly Strings started in Estonia and became famous for their unique blend of folk and contemporary music.
+    They often perform at international festivals and are known for their emotional and poetic lyrics.
+    """}
+]
+
+# ✅ Store Knowledge in ChromaDB (If Not Already Stored)
+existing_data = collection.get()
+if not existing_data["ids"]:
+    for item in knowledge_base:
+        embedding = embedder.encode(item["text"]).tolist()
+        collection.add(documents=[item["text"]], embeddings=[embedding], ids=[item["id"]])
+
+# ✅ Load Llama Model
 llm = Llama.from_pretrained(
     repo_id="krishna195/second_guff",
     filename="unsloth.Q4_K_M.gguf",
 )
 
-# Define the chatbot function
-def chatbot_response(user_input):
-    # System instructions
-    system_prompt = """
-    You are a chatbot specializing in recommending songs by the Estonian folk band **Curly Strings**. 
-    based on this anw the question and give the link of the source 
-    ## 🎵 **Song List**
-    Here are some songs by Curly Strings:
-    1. **Kalakesed**
-    2. **Kus mu süda on ...**
-    3. **Vitsalaul**
-    4. **Viimases jaamas**
-    5. **Salaja**
-    6. **Üle ilma**
-    7. **Šveits**
-    8. **Kallimale**
-    9. **Üksteist peab hoidma**
-    10. **Suuda öelda ei**
-    11. **Annan käe**
-    12. **Tulbid ja Bonsai**
-    13. **Tüdruk Pika Kleidiga**
-    14. **Armasta mind (feat. Vaiko Eplik)**
-    15. **Minu, Pets, Margus ja Priit**
-    16. **Kauges külas**
-    17. **Tule ja jää**
-    18. **Kuutõbine**
-    19. **Omaenese ilus ja veas**
-    20. **Pulmad**
-    21. **Pillimeeste laul**
-    22. **Tehke ruumi!**
-
-    ## 🎤 **Related Artists**
-    If you enjoy Curly Strings, you might also like:
-    - **Trad.Attack!**
-    - **Eesti Raadio laululapsed**
-    - **Körsikud**
-    - **Karl-Erik Taukar**
-    - **Dag**
-    - **Sadamasild**
-    - **Kruuv**
-    - **Smilers**
-    - **Mari Jürjens**
-    - **Terminaator**
+# ✅ Function to Retrieve Relevant Knowledge
+def retrieve_context(query):
+    query_embedding = embedder.encode(query).tolist()
+    results = collection.query(query_embeddings=[query_embedding], n_results=2)
+    retrieved_texts = [doc for doc in results.get("documents", []) if doc]
+    return "\n".join(retrieved_texts) if retrieved_texts else "No relevant data found."
 
-    ---
-    """
+# ✅ Chatbot Function with ChromaDB-RAG
+def chatbot_response(user_input):
+    context = retrieve_context(user_input)  # Retrieve relevant info from ChromaDB
+    
+    messages = [
+        {"role": "system", "content": "Use the knowledge retrieved to answer the user’s question."},
+        {"role": "user", "content": user_input},
+        {"role": "assistant", "content": f"Retrieved Context:\n{context}"},
+    ]
 
-    # Generate response from Llama model
     response = llm.create_chat_completion(
-        messages=[
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_input}
-        ],
+        messages=messages,
         temperature=0.5,
-        max_tokens=1000,  # Increased for better answers
+        max_tokens=500,
         top_p=0.9,
         frequency_penalty=0.8,
     )
 
     return response["choices"][0]["message"]["content"].strip()
 
-# Create Gradio interface
+# ✅ Gradio UI
 iface = gr.Interface(
     fn=chatbot_response,
     inputs=gr.Textbox(placeholder="Ask me about Curly Strings..."),
@@ -78,5 +90,5 @@ iface = gr.Interface(
     theme="compact",
 )
 
-# Launch the Gradio app
+# ✅ Launch App
 iface.launch()