File size: 3,692 Bytes
24e3a93 0296ab6 24e3a93 de0cab5 24e3a93 de0cab5 24e3a93 de0cab5 24e3a93 de0cab5 24e3a93 ccea2c8 24e3a93 de0cab5 24e3a93 0b0e325 24e3a93 |
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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import streamlit as st
import logging
from BanglaRAG.bangla_rag_pipeline import BanglaRAGChain
import warnings
warnings.filterwarnings("ignore")
#lalala
# Default constants for the script
DEFAULT_CHAT_MODEL_ID = "hassanaliemon/bn_rag_llama3-8b"
DEFAULT_EMBED_MODEL_ID = "l3cube-pune/bengali-sentence-similarity-sbert"
DEFAULT_K = 4
DEFAULT_TOP_K = 2
DEFAULT_TOP_P = 0.6
DEFAULT_TEMPERATURE = 0.6
DEFAULT_CHUNK_SIZE = 500
DEFAULT_CHUNK_OVERLAP = 150
DEFAULT_MAX_NEW_TOKENS = 256
DEFAULT_OFFLOAD_DIR = "/tmp"
# Set up logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
# Initialize and load the RAG model
@st.cache_resource(show_spinner=False)
def load_model(chat_model_id, embed_model_id, text_path, k, top_k, top_p, temperature, chunk_size, chunk_overlap, hf_token, max_new_tokens, quantization, offload_dir):
rag_chain = BanglaRAGChain()
rag_chain.load(
chat_model_id=chat_model_id,
embed_model_id=embed_model_id,
text_path=text_path,
k=k,
top_k=top_k,
top_p=top_p,
temperature=temperature,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
hf_token=hf_token,
max_new_tokens=max_new_tokens,
quantization=quantization,
offload_dir=offload_dir, # Pass the offload_dir here
)
return rag_chain
def main():
st.title("Bangla RAG Chatbot")
# Sidebar for model configuration
st.sidebar.header("Model Configuration")
chat_model_id = st.sidebar.text_input("Chat Model ID", DEFAULT_CHAT_MODEL_ID)
embed_model_id = st.sidebar.text_input("Embed Model ID", DEFAULT_EMBED_MODEL_ID)
k = st.sidebar.slider("Number of Documents to Retrieve (k)", 1, 10, DEFAULT_K)
top_k = st.sidebar.slider("Top K", 1, 10, DEFAULT_TOP_K)
top_p = st.sidebar.slider("Top P", 0.0, 1.0, DEFAULT_TOP_P)
temperature = st.sidebar.slider("Temperature", 0.0, 1.0, DEFAULT_TEMPERATURE)
max_new_tokens = st.sidebar.slider("Max New Tokens", 1, 512, DEFAULT_MAX_NEW_TOKENS)
chunk_size = st.sidebar.slider("Chunk Size", 100, 1000, DEFAULT_CHUNK_SIZE)
chunk_overlap = st.sidebar.slider("Chunk Overlap", 0, 500, DEFAULT_CHUNK_OVERLAP)
text_path = st.sidebar.text_input("Text File Path", "text.txt")
quantization = st.sidebar.checkbox("Enable Quantization (4-bit)", value=False)
show_context = st.sidebar.checkbox("Show Retrieved Context", value=False)
offload_dir = st.sidebar.text_input("Offload Directory", DEFAULT_OFFLOAD_DIR) # Default to /tmp
# Load the model with the above configuration
rag_chain = load_model(
chat_model_id=chat_model_id,
embed_model_id=embed_model_id,
text_path=text_path,
k=k,
top_k=top_k,
top_p=top_p,
temperature=temperature,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
hf_token=None, # If you're not using HF API token, set it to None
max_new_tokens=max_new_tokens,
quantization=quantization,
offload_dir=offload_dir, # Pass the offload_dir here
)
st.write("### Enter your question:")
query = st.text_input("আপনার প্রশ্ন")
if st.button("Generate Answer"):
if query:
try:
answer, context = rag_chain(query)
st.write(f"**Answer:** {answer}")
if show_context:
st.write(f"**Context:** {context}")
except Exception as e:
st.error(f"Couldn't generate an answer: {e}")
else:
st.warning("Please enter a query.")
if __name__ == "__main__":
main()
|