Spaces:
Sleeping
Sleeping
from global_vars import translations, t | |
from app import Plugin | |
import streamlit as st | |
import yaml | |
from litellm import completion, embedding | |
import numpy as np | |
from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances, manhattan_distances | |
import os | |
from typing import List, Dict, Any | |
import requests | |
import torch | |
from transformers import AutoTokenizer, AutoModel | |
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' | |
MAX_LENGTH = 512 | |
CHUNK_SIZE = 200 # Nombre de mots par chunk | |
def mean_pooling(model_output, attention_mask): | |
token_embeddings = model_output[0] | |
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() | |
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) | |
# Ajout des traductions spécifiques à ce plugin | |
translations["en"].update({ | |
"rag_plugin_loaded": "RAG LLM Plugin loaded", | |
"rag_enter_text": "Enter RAG text:", | |
"rag_enter_question": "Enter your question:", | |
"rag_button_get_answer": "Get an answer", | |
"rag_success_text_processed": "RAG text processed successfully!", | |
"rag_warning_enter_text": "Please enter RAG text.", | |
"rag_warning_process_text_first": "Please process the RAG text first.", | |
"rag_warning_enter_question": "Please enter a question.", | |
"rag_answer": "Answer:", | |
"rag_citations": "Citations:", | |
"rag_model_provider": "Model Provider", | |
"rag_llm_model": "LLM Model", | |
"rag_embedder_model": "Embedding Model", | |
"rag_similarity_method": "Similarity Method", | |
"rag_llm_sys_prompt": "System prompt for LLM", | |
"rag_chunk_size": "Chunk size", | |
"rag_top_k_chunks": "Number of chunks to use", | |
"rag_default_sys_prompt": "You are an AI assistant. Your task is to analyze the provided context and answer questions based ONLY on this context. If the information is not in the context, clearly state that.", | |
"rag_error_fetching_models_ollama": "Error fetching Ollama models: ", | |
"rag_error_calling_llm": "Error calling LLM: ", | |
"rag_processing" : "Processing...", | |
}) | |
translations["fr"].update({ | |
"rag_plugin_loaded": "Plugin RAG LLM chargé", | |
"rag_enter_text": "Entrez le texte RAG :", | |
"rag_enter_question": "Entrez votre question :", | |
"rag_button_get_answer": "Obtenir une réponse", | |
"rag_success_text_processed": "Texte RAG traité avec succès!", | |
"rag_warning_enter_text": "Veuillez entrer du texte RAG.", | |
"rag_warning_process_text_first": "Veuillez d'abord traiter le texte RAG.", | |
"rag_warning_enter_question": "Veuillez entrer une question.", | |
"rag_answer": "Réponse :", | |
"rag_citations": "Citations :", | |
"rag_model_provider": "Fournisseur de modèle", | |
"rag_llm_model": "Modèle LLM", | |
"rag_embedder_model": "Modèle d'embedding", | |
"rag_similarity_method": "Méthode de similarité", | |
"rag_llm_sys_prompt": "Prompt système pour le LLM", | |
"rag_chunk_size": "Taille des chunks", | |
"rag_top_k_chunks": "Nombre de chunks à utiliser", | |
"rag_default_sys_prompt": "Tu es un assistant IA. Ta tâche est d'analyser le contexte fourni et de répondre aux questions en te basant UNIQUEMENT sur ce contexte. Si l'information n'est pas dans le contexte, dis-le clairement.", | |
"rag_error_fetching_models_ollama": "Erreur lors de la récupération des modèles Ollama : ", | |
"rag_error_calling_llm": "Erreur lors de l'appel au LLM : ", | |
"rag_processing" : "En cours de traitement...", | |
}) | |
class RagllmPlugin(Plugin): | |
def __init__(self, name: str, plugin_manager): | |
super().__init__(name, plugin_manager) | |
self.config = self.load_llm_config() | |
self.embeddings = None | |
self.chunks = None | |
def load_llm_config(self) -> Dict: | |
with open('.llm-config.yml', 'r') as file: | |
return yaml.safe_load(file) | |
def get_tabs(self): | |
return [{"name": "RAG", "plugin": "ragllm"}] | |
def get_config_fields(self): | |
return { | |
"provider": { | |
"type": "select", | |
"label": t("rag_model_provider"), | |
"options": [("ollama", "Ollama"), ("groq", "Groq")], | |
"default": "ollama" | |
}, | |
"llm_model": { | |
"type": "select", | |
"label": t("rag_llm_model"), | |
"options": [("none", "À charger...")], | |
"default": "ollama/qwen2" | |
}, | |
"embedder": { | |
"type": "select", | |
"label": t("rag_embedder_model"), | |
"options": [ | |
("sentence-transformers/all-MiniLM-L6-v2", "all-MiniLM-L6-v2"), | |
("nomic-ai/nomic-embed-text-v1.5", "nomic-embed-text-v1.5") | |
], | |
"default": "sentence-transformers/all-MiniLM-L6-v2" | |
}, | |
"similarity_method": { | |
"type": "select", | |
"label": t("rag_similarity_method"), | |
"options": [ | |
("cosine", "Cosinus"), | |
("euclidean", "Distance euclidienne"), | |
("manhattan", "Distance de Manhattan") | |
], | |
"default": "cosine" | |
}, | |
"llm_sys_prompt": { | |
"type": "textarea", | |
"label": t("rag_llm_sys_prompt"), | |
"default": t("rag_default_sys_prompt") | |
}, | |
"chunk_size": { | |
"type": "number", | |
"label": t("rag_chunk_size"), | |
"default": 200 | |
}, | |
"top_k": { | |
"type": "number", | |
"label": t("rag_top_k_chunks"), | |
"default": 3 | |
} | |
} | |
def get_config_ui(self, config): | |
updated_config = {} | |
for field, params in self.get_config_fields().items(): | |
if params['type'] == 'select': | |
if field == 'llm_model': | |
provider = config.get('provider', 'ollama') | |
models = self.get_available_models(provider) | |
try: | |
default_index = models.index(config.get(field, params['default'])) | |
except ValueError: | |
default_index = 0 | |
updated_config[field] = st.selectbox( | |
params['label'], | |
options=models, | |
index=default_index | |
) | |
else: | |
options_list = [option[0] for option in params['options']] | |
try: | |
default_index = options_list.index(config.get(field, params['default'])) | |
except ValueError: | |
default_index = 0 | |
updated_config[field] = st.selectbox( | |
params['label'], | |
options=options_list, | |
format_func=lambda x: dict(params['options'])[x], | |
index=default_index | |
) | |
elif params['type'] == 'textarea': | |
updated_config[field] = st.text_area( | |
params['label'], | |
value=config.get(field, params['default']) | |
) | |
elif params['type'] == 'number': | |
updated_config[field] = st.number_input( | |
params['label'], | |
value=int(config.get(field, params['default'])), | |
step=1 | |
) | |
else: | |
updated_config[field] = st.text_input( | |
params['label'], | |
value=config.get(field, params['default']) | |
) | |
return updated_config | |
def get_sidebar_config_ui(self, config: Dict[str, Any]) -> Dict[str, Any]: | |
available_models = self.get_available_models('ollama') + self.get_available_models('groq') | |
default_model = config.get('llm_model', available_models[0] if available_models else None) | |
selected_model = st.sidebar.selectbox( | |
t("rag_llm_model"), | |
options=available_models, | |
index=available_models.index(default_model) if default_model in available_models else 0, | |
key="ragllm_llm_model" | |
) | |
return {"llm_model": selected_model} | |
def get_available_models(self, provider: str) -> List[str]: | |
if provider == 'ollama': | |
try: | |
response = requests.get("http://localhost:11434/api/tags") | |
models = response.json()["models"] | |
return [f"ollama/{model['name']}" for model in models] + ["ollama/qwen2"] | |
except Exception as e: | |
st.error(f"{t('rag_error_fetching_models_ollama')}{str(e)}") | |
return ["ollama/qwen2"] | |
elif provider == 'groq': | |
return ["groq/llama3-70b-8192", "groq/mixtral-8x7b-32768"] | |
else: | |
return ["none"] | |
def process_rag_text(self, rag_text: str, chunk_size: int, embedder): | |
rag_text = rag_text.replace('\\n', ' ').replace('\\\'', "'") | |
mots = rag_text.split() | |
self.chunks = [' '.join(mots[i:i+chunk_size]) for i in range(0, len(mots), chunk_size)] | |
self.embeddings = np.vstack([self.get_embedding(c, embedder) for c in self.chunks]) | |
def get_embedding(self, text: str, model: str) -> np.ndarray: | |
tokenizer = AutoTokenizer.from_pretrained(model) | |
model = AutoModel.from_pretrained(model, trust_remote_code=True).to(DEVICE) | |
inputs = tokenizer(text, padding=True, truncation=True, max_length=MAX_LENGTH, return_tensors="pt").to(DEVICE) | |
with torch.no_grad(): | |
model_output = model(**inputs) | |
return mean_pooling(model_output, inputs['attention_mask']).cpu().numpy() | |
def calculate_similarity(self, query_embedding: np.ndarray, method: str) -> np.ndarray: | |
if method == 'cosine': | |
return cosine_similarity(query_embedding.reshape(1, -1), self.embeddings)[0] | |
elif method == 'euclidean': | |
return -euclidean_distances(query_embedding.reshape(1, -1), self.embeddings)[0] | |
elif method == 'manhattan': | |
return -manhattan_distances(query_embedding.reshape(1, -1), self.embeddings)[0] | |
else: | |
raise ValueError("Méthode de similarité non reconnue") | |
def get_context(self, query: str, config: Dict[str, Any]) -> tuple: | |
query_embedding = self.get_embedding(query, config['ragllm']['embedder']) | |
similarities = self.calculate_similarity(query_embedding, config['ragllm']['similarity_method']) | |
top_indices = np.argsort(similarities)[-config['ragllm']['top_k']:][::-1] | |
context = "\n\n".join([self.chunks[i] for i in top_indices]) | |
return context, [self.chunks[i] for i in top_indices] | |
def call_llm(self, prompt: str, sysprompt: str) -> str: | |
try: | |
llm_model = st.session_state.ragllm_llm_model | |
#print(f"---------------------------------------\nCalling LLM {llm_model} \n with sysprompt {sysprompt} \n and prompt {prompt} \n and context len of {len(context)}") | |
messages = [ | |
{"role": "system", "content": sysprompt}, | |
{"role": "user", "content": prompt} | |
] | |
response = completion(model=llm_model, messages=messages) | |
return response['choices'][0]['message']['content'] | |
except Exception as e: | |
return f"{t('rag_error_calling_llm')}{str(e)}" | |
def free_llm(self): | |
try: | |
llm_model = st.session_state.ragllm_llm_model | |
if llm_model.startswith("ollama/"): | |
ollama_model = llm_model.split("/")[1] | |
response = requests.post( | |
"http://localhost:11434/api/generate", | |
json={ | |
"model": ollama_model, | |
"prompt": "bye", | |
"keep_alive": 0 | |
} | |
) | |
return response.json()['response'] | |
except Exception as e: | |
return f"{t('rag_error_calling_llm')}{str(e)}" | |
def process_with_llm(self, prompt: str, sysprompt: str, context: str) -> str: | |
return self.call_llm(f"Contexte : {context}\n\nQuestion : {prompt}", sysprompt) | |
def run(self, config): | |
st.write(t("rag_plugin_loaded")) | |
# Initialiser rag_text avec la valeur de session_state si elle existe, sinon utiliser une chaîne vide | |
if 'rag_text' not in st.session_state: | |
st.session_state.rag_text = "" | |
if 'rag_question' not in st.session_state: | |
st.session_state.rag_question = "Question" | |
rag_text = st.text_area(t("rag_enter_text"), height=200, value=st.session_state.rag_text, key="rag_text_key") | |
user_prompt = st.text_area(t("rag_enter_question"), value=st.session_state.rag_question, key="rag_prompt_key") | |
st.session_state.rag_text = rag_text # Mettre à jour la valeur dans session_state | |
st.session_state.rag_question = user_prompt | |
if st.button(t("rag_button_get_answer"), key="get_answer_button"): | |
with st.spinner(t("rag_processing")): | |
if rag_text: | |
self.process_rag_text(rag_text, config['ragllm']['chunk_size'], config['ragllm']['embedder']) | |
st.success(t("rag_success_text_processed")) | |
else: | |
st.warning(t("rag_warning_enter_text")) | |
if user_prompt and self.embeddings is not None: | |
context, citations = self.get_context(user_prompt, config) | |
response = self.process_with_llm(user_prompt, config['ragllm']['llm_sys_prompt'], context) | |
st.write(t("rag_answer")) | |
st.write(response) | |
st.write(t("rag_citations")) | |
for i, citation in enumerate(citations, 1): | |
st.write(f"{i}. {citation[:100]}...") | |
elif self.embeddings is None: | |
st.warning(t("rag_warning_process_text_first")) | |
else: | |
st.warning(t("rag_warning_enter_question")) | |