File size: 3,297 Bytes
b70b72a |
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 |
from pathlib import Path
# import base64
import gradio as gr
from src.rag import Rag
from src.amodel import ModelType
STORE_DIR = "./files/rag_app" # Le répertoire de la base
MAX_DOCS = 4 # Le nombre max de documents dans la base
def main():
# Création du rag
rag:Rag = Rag(ModelType.MTHUGGINGFACE, store_dir=STORE_DIR)
# Reset de la base à chaque démarrage du serveur
rag.reset_store()
# UI
with gr.Blocks() as demo:
def upload_file(file_path):
# Récupérer la liste des collections du store
names = list(rag.emb_store.get_collection_names())
# Supprimer la première s'il y en a déjà 4 (le max)
if len(names) == MAX_DOCS:
rag.delete_collection(names[0])
# Ajouter le pdf
name:str = Path(file_path).name
rag.add_pdf_to_store(file_name=file_path, collection_name=name)
# Mettre à jour la liste des noms de collections
names = list(rag.emb_store.get_collection_names())
combo.choices = names
return gr.update(choices=names, value=name, interactive=True), gr.update(value="")
def ask_rag(question:str, choice:str):
col_name:str = choice
if col_name == None:
return "Aucun pdf actif, veuillez en uploader un !"
prompt, resp, sources, ids = rag.ask_rag(question, col_name)
# print("choice:" + col_name, "RAG:" + resp)
return resp
def on_temperature_change(temp):
rag.set_temperature(temp)
with gr.Tab("RAG"):
# with gr.Row():
gr.Image("./files/drane.jpg", height=100, show_download_button=False, show_fullscreen_button=False, show_label=False)
# Le combo qui permet le choix du pdf
names:list[str] = rag.emb_store.get_collection_names()
combo = gr.Dropdown(names, label="PDFs", multiselect=False)
# Le button qui permet d'uploader un pdf
upload_button = gr.UploadButton("Clique pour uploader un pdf", file_types=[".pdf"], file_count="single")
# La zone où on pose une question au RAG
ask_input = gr.Text(label="Pose une question à ton pdf")
# La réponse du RAG (Markdown pour afficher les formules .tex)
rag_output = gr.Markdown(label="Réponse")
with gr.Tab("Réglages"):
gr.Markdown("## Modèles:")
gr.Markdown("- " + rag.get_llm_name())
gr.Markdown("- " + rag.get_feature_name())
temperature_slider = gr.Slider(minimum=0,
maximum=1.0,
value=0.5,
step=0.1,
label="Température")
# Réponses aux évènements
upload_button.upload(fn=upload_file, inputs=upload_button, outputs=[combo, rag_output], show_progress="full")
ask_input.submit(fn=ask_rag, inputs=[ask_input, combo], outputs=rag_output)
temperature_slider.change(fn=on_temperature_change, inputs=temperature_slider)
# demo.launch(allowed_paths=["./file/"], share=True)
demo.launch(allowed_paths=["./file/"])
if __name__ == "__main__":
main() |