Spaces:
Sleeping
Sleeping
| 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.") | |