Spaces:
Sleeping
Sleeping
File size: 2,831 Bytes
b93b2dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import streamlit as st
from utils.vector_base import KnowledgeBase
from utils.embedding import Embeddings
from utils.llm import LLM
from config import config
import json
def get_emdedding_model():
return Embeddings()
def get_llm(url, api_key):
return LLM(url, api_key)
def get_metadata(path):
titles, texts = [], []
with open(path, 'rb') as file:
metadata = json.load(file)
for data in metadata:
titles.append(data['title'])
texts.append(data['text'])
return texts, titles
def combine_docs(indexes, texts):
result = ""
for i, index in enumerate(indexes):
result += " [" + str(i + 1) + "] " + texts[index]
return result
def create_prompt(query, docs):
system_prompt = f"""You are a language model integrated into a search and
generation system based on relevant documents (RAG system).
Your task is to provide answers to the user's queries based on the provided
documents. Respond only based on the provided documents. Do not make up
information that is not in the sources. If you use data from a document,
indicate the document number in square brackets. For example: "This term
means such-and-such [1]." If there is no information in the documents,
politely explain that the information is not available. Do not alter the
content of the sources, convey the information accurately.
User query: {query}. Documents: {docs}
"""
return system_prompt
st.title("PaperRAG")
st.write("RAG system for scientific papers with selectable search types")
query = st.text_input("Enter your query", "")
search_types = st.multiselect(
"Search Types", options=["Vector", "BM25"], default=["Vector", "BM25"]
)
llm_url = st.text_input("LLM URL", "", placeholder="Enter LLM ENDPOINT")
llm_api_key = st.text_input("LLM API Key", "", placeholder="Enter LLM API Key", type="password")
if st.button("Search"):
if query and llm_url and llm_api_key:
model = get_emdedding_model()
llm = get_llm(llm_url, llm_api_key)
texts, titles = get_metadata(config.PATH_METADATA)
embedding = model.get_query_embedding(query)
knowledge_base = KnowledgeBase(config.PATH_FAISS, config.PATH_PREPROCESSING_TEXT)
vector_search = []
bm25_search = []
if "Vector" in search_types:
vector_search = knowledge_base.search_by_embedding(embedding, 5)[0].tolist()
if "BM25" in search_types:
bm25_search = knowledge_base.search_by_BM25(query, 5)
docs = combine_docs(vector_search + bm25_search, texts)
prompt = create_prompt(query, docs)
response = llm.generate_response(prompt)
st.subheader("Response")
st.write(response)
else:
st.error("Please fill in all the required fields.")
|