Update app.py

app.py

@@ -28,8 +28,8 @@ async def generate_keywords(query):
     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,
+        model="Qwen/Qwen2.5-72B-Instruct",
+        prompt=prompt,
         max_new_tokens=50,
         temperature=0.3,
         top_p=0.9
@@ -38,15 +38,6 @@ async def generate_keywords(query):
     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(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)
     
@@ -55,15 +46,12 @@ async def generate_ai_response(query, 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}"}
-    ]
+    prompt = f"Jesteś asystentem prawniczym. Odpowiedz na poniższe pytanie na podstawie podanego kontekstu prawnego.\n\nKontekst: {context}\n\nPytanie: {query}\n\nOdpowiedź:"
 
     response = ""
     async for token in client.text_generation(
-        "Qwen/Qwen2.5-72B-Instruct",
-        messages,
+        model="Qwen/Qwen2.5-72B-Instruct",
+        prompt=prompt,
         max_new_tokens=2048,
         temperature=0.5,
         top_p=0.7,
@@ -72,6 +60,16 @@ async def generate_ai_response(query, relevant_chunks):
         response += token
         yield token
 
+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(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]
+
+
 def main():
     st.title("Chatbot Prawny z AI")