Update views/rag_med.py

views/rag_med.py

@@ -9,24 +9,6 @@ from sentence_transformers import SentenceTransformer
 import pickle
 from prompt_template import prompt
 
-# Função para encontrar o arquivo
-def find_file(filename):
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    root_dir = os.path.dirname(current_dir)
-    
-    possible_paths = [
-        os.path.join(current_dir, filename),
-        os.path.join(root_dir, filename),
-        os.path.join("/content", filename),  # Para ambientes como Colab
-        os.path.join("/home/user/app/", filename)  # Caminho comum no Hugging Face
-    ]
-    
-    for path in possible_paths:
-        if os.path.exists(path):
-            return path
-    
-    raise FileNotFoundError(f"Não foi possível encontrar o arquivo: {filename}")
-
 # Initialize the messages list in the session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -54,39 +36,41 @@ with st.sidebar:
         st.session_state.messages = []
         st.rerun()
 
-# Carregando o índice FAISS e os dados armazenados
-@st.cache_resource
-def load_faiss_index():
-    index_file = find_file("faiss_index.bin")
-    data_file = find_file("stored_data.pkl")
-    
-    index = faiss.read_index(index_file)
-    with open(data_file, 'rb') as f:
-        stored_data = pickle.load(f)
-    return index, stored_data
-
-try:
-    index, stored_data = load_faiss_index()
-except Exception as e:
-    st.error(f"Erro ao carregar o índice FAISS: {e}")
-    st.stop()
-
-# Inicialize o modelo de embedding
 model = SentenceTransformer("all-MiniLM-L6-v2")
 
-def search(query, k=3):
+# índice FAISS
+json_file = "scr/faiss/chunks .json"
+embeddings_file = "scr/faiss/embeddings.npy"
+index_file = "scr/faiss/faiss_index.index"
+text_chunks_file = "scr/faiss/text_chunks.npy"
+
+with open(json_file, "r") as file:
+    chunks = json.load(file)
+
+embeddings = np.load(embeddings_file)
+index = faiss.read_index(index_file)
+text_chunks = np.load(text_chunks_file)
+
+with st.sidebar:
+    openai_api_key = st.text_input("OpenAI API Key",
+                                   key="chatbot_api_key",
+                                   type="password")
+    st.markdown(
+        "[Pegue aqui sua chave OpenAI API](https://platform.openai.com/account/api-keys)"
+    )
+if st.sidebar.button("Limpar Conversa"):
+    st.session_state.messages = []
+    st.rerun()
+
+
+def retrieve(query, k=3):
+    """
+    Retrieves the most similar document chunks to a given query.
+    """
     query_embedding = model.encode([query])
-    query_embedding = query_embedding.astype(np.float32).reshape(1, -1)
     distances, indices = index.search(query_embedding, k)
-    
-    results = []
-    for i, idx in enumerate(indices[0]):
-        if idx != -1 and idx < len(stored_data):
-            result = stored_data[idx]
-            result['score'] = float(1 - distances[0][i]/2)  # Convertendo distância para similaridade
-            results.append(result)
-    
-    return results
+    return [chunks[idx] for idx in indices[0]]
+
 
 def generate_response(query, context):
     """
@@ -105,8 +89,10 @@ def generate_response(query, context):
         model="gpt-4o-mini",
         messages=[
             {
-                "role": "system",
-                "content": "Você é um assistente especializado em análise de encaminhamentos médicos.",
+                "role":
+                "system",
+                "content":
+                "Você é um assistente especializado em análise de encaminhamentos médicos.",
             },
             {
                 "role": "user",
@@ -116,19 +102,19 @@ def generate_response(query, context):
     )
     return response.choices[0].message.content
 
+
 co1, co2, co3 = st.columns([1.5, 0.4, 3])
 with co1:
     st.write('')
 with co2:
-    st.write('')
-    #st.image('icons/icon.svg', width=80)
+    st.image('icons/icon.svg', width=80)
 
 with co3:
     st.title("Encaminhamento Médico")
 
 col1, col2, col3, col4, col5 = st.columns([1, 3, 0.5, 3, 1])
 
-
+fem = 'icons/iconF.svg'
 with col1:
     st.write('')
 
@@ -139,8 +125,7 @@ with col2:
     cl1, cl2, cl3 = st.columns([2, 1, 2])
     with cl1:
         idade = st.text_input("###### IDADE", value="")
-        sexo_opcao = st.radio("Sexo", ("Masculino", "Feminino"),
-              index=0)
+        sexo_opcao = st.radio("Sexo", ("Masculino", "Feminino"), index=0)
         sexo = (sexo_opcao == "Masculino")
     with cl2:
         st.write("")
@@ -253,7 +238,6 @@ with col4:
                     mas adaptando o conteúdo para ser mais fluido e coeso.
                     """
         st.text_area("Prompt gerado:", user_prompt, height=600)
-      
         if not openai_api_key:
             st.info("Insira sua chave API OpenAI para continuar.")
             st.stop()
@@ -291,9 +275,8 @@ if prompt := st.chat_input("Digite sua pergunta"):
         st.markdown(prompt)
 
     # Processa a query e gera a resposta
-    search_results = search(prompt)
-    context = "\n".join([f"Text: {result['text'][:200]}..." for result in search_results])
-    response = generate_response(prompt, context)
+    retrieved_context = retrieve(prompt)
+    response = generate_response(prompt, retrieved_context)
 
     # Adiciona a resposta do bot
     st.session_state.messages.append({
@@ -303,4 +286,4 @@ if prompt := st.chat_input("Digite sua pergunta"):
     with st.chat_message("assistant"):
         st.markdown(response)
 with col5:
-    st.write('')
+    st.write('')