Spaces:
Sleeping
Sleeping
File size: 2,931 Bytes
46e28ab 3ad9a49 46e28ab f3576a5 46e28ab f3576a5 4e65999 46e28ab 4e65999 f3576a5 114ce4a 3ad9a49 114ce4a f3576a5 114ce4a 4e65999 3ad9a49 4e65999 46e28ab 4e65999 |
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 |
# AI assistant with a RAG system to query information from the CAMELS cosmological simulations using Langchain and deployed with Gradio
# Author: Pablo Villanueva Domingo
from rag import RAG, load_docs
from langchain_community.embeddings import HuggingFaceInstructEmbeddings
from langchain_mistralai import ChatMistralAI
from langchain_core.rate_limiters import InMemoryRateLimiter
import gradio as gr
# Define a limiter to avoid rate limit issues with MistralAI
rate_limiter = InMemoryRateLimiter(
requests_per_second=0.1, # <-- MistralAI free. We can only make a request once every second
check_every_n_seconds=0.01, # Wake up every 100 ms to check whether allowed to make a request,
max_bucket_size=10, # Controls the maximum burst size.
)
# Load the documentation
docs = load_docs()
print("Pages loaded:",len(docs))
# LLM model
llm = ChatMistralAI(model="mistral-large-latest", rate_limiter=rate_limiter)
# Embeddings
embed_model = "sentence-transformers/multi-qa-distilbert-cos-v1"
# embed_model = "nvidia/NV-Embed-v2"
embeddings = HuggingFaceInstructEmbeddings(model_name=embed_model)
# RAG chain
rag_chain = RAG(llm, docs, embeddings)
# Function to handle prompt and query the RAG chain
def handle_prompt(message, history):
try:
# Stream output
out=""
for chunk in rag_chain.stream(message):
out += chunk
yield out
except:
raise gr.Error("Requests rate limit exceeded")
if __name__=="__main__":
# Predefined messages and examples
description = "AI powered assistant which answers any question related to the [CAMELS simulations](https://www.camel-simulations.org/)."
greetingsmessage = "Hi, I'm the CAMELS DocBot, I'm here to assist you with any question related to the CAMELS simulations."
example_questions = [
"How can I read a halo file?",
"Which simulation suites are included in CAMELS?",
"Which are the largest volumes in CAMELS simulations?",
"Write a complete snippet of code getting the power spectrum of a simulation"
]
# Define customized Gradio chatbot
chatbot = gr.Chatbot([{"role":"assistant", "content":greetingsmessage}],
type="messages",
avatar_images=["ims/userpic.png","ims/camelslogo.jpg"],
height="60vh")
# Define Gradio interface
demo = gr.ChatInterface(handle_prompt,
type="messages",
title="CAMELS DocBot",
fill_height=True,
examples=example_questions,
theme=gr.themes.Soft(),
description=description,
#cache_examples=False,
chatbot=chatbot)
demo.launch() |