Update app.py

app.py

@@ -3,6 +3,13 @@ import json
 import os
 from sentence_transformers import SentenceTransformer, util
 import torch
+from huggingface_hub import InferenceClient
+import asyncio
+
+# Load the Hugging Face token from environment variable
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is not set. Please set it before running the application.")
 
 # Load the processed legal code data
 @st.cache_resource
@@ -15,17 +22,58 @@ def load_data(file_path):
 def load_model():
     return SentenceTransformer('distiluse-base-multilingual-cased-v1')
 
-def search_relevant_chunks(query, chunks, model, top_k=3):
-    query_embedding = model.encode(query, convert_to_tensor=True)
+async def generate_keywords(query):
+    client = InferenceClient(token=HF_TOKEN)
+    
+    prompt = f"Na podstawie poniższego pytania, wygeneruj 3-5 słów kluczowych, które najlepiej opisują główne tematy i koncepcje prawne zawarte w pytaniu. Podaj tylko słowa kluczowe, oddzielone przecinkami.\n\nPytanie: {query}\n\nSłowa kluczowe:"
+    
+    response = await client.text_generation(
+        "Qwen/Qwen2.5-72B-Instruct",
+        prompt,
+        max_new_tokens=50,
+        temperature=0.3,
+        top_p=0.9
+    )
+    
+    keywords = [keyword.strip() for keyword in response.split(',')]
+    return keywords
+
+def search_relevant_chunks(keywords, chunks, model, top_k=3):
+    keyword_embedding = model.encode(keywords, convert_to_tensor=True)
     chunk_embeddings = model.encode([chunk['text'] for chunk in chunks], convert_to_tensor=True)
     
-    cos_scores = util.pytorch_cos_sim(query_embedding, chunk_embeddings)[0]
-    top_results = torch.topk(cos_scores, k=top_k)
+    cos_scores = util.pytorch_cos_sim(keyword_embedding, chunk_embeddings)
+    top_results = torch.topk(cos_scores.mean(dim=0), k=top_k)
     
     return [chunks[idx] for idx in top_results.indices]
 
+async def generate_ai_response(query, relevant_chunks):
+    client = InferenceClient(token=HF_TOKEN)
+    
+    context = "Kontekst prawny:\n\n"
+    for chunk in relevant_chunks:
+        context += f"{chunk['metadata']['nazwa']} - Artykuł {chunk['metadata']['article']}:\n"
+        context += f"{chunk['text']}\n\n"
+
+    messages = [
+        {"role": "system", "content": "Jesteś asystentem prawniczym. Odpowiadaj na pytania na podstawie podanego kontekstu prawnego."},
+        {"role": "user", "content": f"Kontekst: {context}\n\nPytanie: {query}"}
+    ]
+
+    response = ""
+    async for token in client.text_generation(
+        "Qwen/Qwen2.5-72B-Instruct",
+        messages,
+        max_new_tokens=2048,
+        temperature=0.5,
+        top_p=0.7,
+        stream=True
+    ):
+        response += token
+        yield token
+
 def main():
-    st.title("Chatbot Prawny")
+    st.title("Chatbot Prawny z AI")
 
     # Load data and model
     data_file = "processed_kodeksy.json"
@@ -51,19 +99,21 @@ def main():
         with st.chat_message("user"):
             st.markdown(prompt)
 
-        # Search for relevant chunks
-        relevant_chunks = search_relevant_chunks(prompt, chunks, model)
-
-        # Generate response
-        response = "Oto co znalazłem w kodeksie:\n\n"
-        for chunk in relevant_chunks:
-            response += f"**{chunk['metadata']['nazwa']} - Artykuł {chunk['metadata']['article']}**\n"
-            response += f"{chunk['text']}\n\n"
+        # Generate keywords and search for relevant chunks
+        with st.spinner("Analizuję pytanie i szukam odpowiednich informacji..."):
+            keywords = asyncio.run(generate_keywords(prompt))
+            relevant_chunks = search_relevant_chunks(keywords, chunks, model)
 
-        # Display assistant response
+        # Generate AI response
         with st.chat_message("assistant"):
-            st.markdown(response)
-        st.session_state.messages.append({"role": "assistant", "content": response})
+            message_placeholder = st.empty()
+            full_response = ""
+            for chunk in asyncio.run(generate_ai_response(prompt, relevant_chunks)):
+                full_response += chunk
+                message_placeholder.markdown(full_response + "▌")
+            message_placeholder.markdown(full_response)
+        
+        st.session_state.messages.append({"role": "assistant", "content": full_response})
 
     # Sidebar for additional options
     with st.sidebar: