fix

rag.py

@@ -1,7 +1,9 @@
 # %%
+
 import os
 import json
 
+
 import torch
 import faiss
 import numpy as np
@@ -13,8 +15,10 @@ from transformers import (
     AutoTokenizer,
 )
 
+
 HF_TOKEN = os.environ["hf_token"]
 
+
 SYSTEM_PROMPT = """You are a helpful question answering assistant. You will be given a context and a question. You need to provide the answer to the question based on the context. Answer briefly, based on the context. Only output the answer, and nothing else. Here is an example:
 
 >> Context
@@ -40,12 +44,14 @@ USER_PROMPT = """
 """
 
 
+
 def load_embedder(model_path: str, device: str) -> SentenceTransformer:
     embedder = SentenceTransformer(model_path)
     embedder.to(device)
     return embedder
 
 
+
 def load_contexts(context_file: str) -> list[str]:
     contexts = []
     with open(context_file, "r") as f_in:
@@ -56,10 +62,12 @@ def load_contexts(context_file: str) -> list[str]:
     return contexts
 
 
+
 def load_index(index_file: str) -> faiss.Index:
     return faiss.read_index(index_file)
 
 
+
 def load_reader(model_path: str, device: str) -> TextGenerationPipeline:
     model = AutoModelForCausalLM.from_pretrained(model_path, token=HF_TOKEN)
 
@@ -78,6 +86,7 @@ def load_reader(model_path: str, device: str) -> TextGenerationPipeline:
     return reader
 
 
+
 def construct_prompt(contexts: list[str], question: str) -> list[dict]:
     return [
         {"role": "system", "content": SYSTEM_PROMPT},
@@ -90,6 +99,7 @@ def construct_prompt(contexts: list[str], question: str) -> list[dict]:
     ]
 
 
+
 def load_all(
     embedder_path: str,
     context_file: str,
@@ -110,6 +120,7 @@ def load_all(
     }
 
 
+
 def run_query(
     question: str,
     embedder: SentenceTransformer,