Update app.py

app.py

@@ -4,7 +4,6 @@ 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")
@@ -22,12 +21,12 @@ def load_data(file_path):
 def load_model():
     return SentenceTransformer('distiluse-base-multilingual-cased-v1')
 
-async def generate_keywords(query):
+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(
+    response = client.text_generation(
         model="Qwen/Qwen2.5-72B-Instruct",
         prompt=prompt,
         max_new_tokens=50,
@@ -38,7 +37,16 @@ async def generate_keywords(query):
     keywords = [keyword.strip() for keyword in response.split(',')]
     return keywords
 
-async def generate_ai_response(query, relevant_chunks):
+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 generate_ai_response(query, relevant_chunks):
     client = InferenceClient(token=HF_TOKEN)
     
     context = "Kontekst prawny:\n\n"
@@ -48,27 +56,15 @@ async def generate_ai_response(query, relevant_chunks):
 
     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(
+    response = client.text_generation(
         model="Qwen/Qwen2.5-72B-Instruct",
         prompt=prompt,
         max_new_tokens=2048,
         temperature=0.5,
-        top_p=0.7,
-        stream=True
-    ):
-        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)
+        top_p=0.7
+    )
     
-    return [chunks[idx] for idx in top_results.indices]
-
+    return response
 
 def main():
     st.title("Chatbot Prawny z AI")
@@ -99,16 +95,13 @@ def main():
 
         # Generate keywords and search for relevant chunks
         with st.spinner("Analizuję pytanie i szukam odpowiednich informacji..."):
-            keywords = asyncio.run(generate_keywords(prompt))
+            keywords = generate_keywords(prompt)
             relevant_chunks = search_relevant_chunks(keywords, chunks, model)
 
         # Generate AI response
         with st.chat_message("assistant"):
             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 + "▌")
+            full_response = generate_ai_response(prompt, relevant_chunks)
             message_placeholder.markdown(full_response)
         
         st.session_state.messages.append({"role": "assistant", "content": full_response})