#MisterAI/Docker_Ollama
#app.py_03
#https://huggingface.co/spaces/MisterAI/Docker_Ollama/
#Ajouter Historique

import logging
import requests
from pydantic import BaseModel
from langchain_community.llms import Ollama
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
import gradio as gr
import threading
import subprocess
from bs4 import BeautifulSoup

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Cache pour stocker les modèles déjà chargés
loaded_models = {}

# Variable pour suivre l'état du bouton "Stop"
stop_flag = False

def get_model_list():
    url = "https://ollama.com/search"
    response = requests.get(url)

    # Vérifier si la requête a réussi
    if response.status_code == 200:
        # Utiliser BeautifulSoup pour analyser le HTML
        soup = BeautifulSoup(response.text, 'html.parser')
        model_list = []

        # Trouver tous les éléments de modèle
        model_elements = soup.find_all('li', {'x-test-model': True})

        for model_element in model_elements:
            model_name = model_element.find('span', {'x-test-search-response-title': True}).text.strip()
            size_elements = model_element.find_all('span', {'x-test-size': True})

            # Filtrer les modèles par taille
            for size_element in size_elements:
                size = size_element.text.strip().lower()  # Convertir en minuscules
                if 'x' in size:
                    # Exclure les modèles avec des tailles de type nXm ou nXb
                    continue
                elif size.endswith('m'):
                    # Tous les modèles en millions sont acceptés
                    model_list.append(f"{model_name}:{size}")
                elif size.endswith('b'):
                    # Convertir les modèles en milliards en milliards
                    size_value = float(size[:-1])
                    if size_value <= 10:  # Filtrer les modèles <= 10 milliards de paramètres
                        model_list.append(f"{model_name}:{size}")

        return model_list
    else:
        logger.error(f"Erreur lors de la récupération de la liste des modèles : {response.status_code} - {response.text}")
        return []

def get_llm(model_name):
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    return Ollama(model=model_name, callback_manager=callback_manager)

class InputData(BaseModel):
    model_name: str
    input: str
    max_tokens: int = 256
    temperature: float = 0.7

def pull_model(model_name):
    try:
        # Exécuter la commande pour tirer le modèle
        subprocess.run(["ollama", "pull", model_name], check=True)
        logger.info(f"Model {model_name} pulled successfully.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to pull model {model_name}: {e}")
        raise

def check_and_load_model(model_name):
    # Vérifier si le modèle est déjà chargé
    if model_name in loaded_models:
        logger.info(f"Model {model_name} is already loaded.")
        return loaded_models[model_name]
    else:
        logger.info(f"Loading model {model_name}...")
        # Tirer le modèle si nécessaire
        pull_model(model_name)
        llm = get_llm(model_name)
        loaded_models[model_name] = llm
        return llm

# Interface Gradio
def gradio_interface(model_name, input, max_tokens, temperature, history):
    global stop_flag
    stop_flag = False

    # Ajouter la requête de l'utilisateur à l'historique une seule fois
    history.append((input, ""))
    yield history, history

    def worker():
        llm = check_and_load_model(model_name)
        response = ""
        for token in llm.stream(input, max_tokens=max_tokens, temperature=temperature):
            if stop_flag:
                break
            response += token
            history[-1] = (input, response)
            yield history, history

    # Utiliser un thread pour gérer le streaming
    for result in worker():
        if stop_flag:
            break
        yield result

model_list = get_model_list()

demo = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Dropdown(model_list, label="Select Model", value="mistral:7b"),
        gr.Textbox(label="Input"),
        gr.Slider(minimum=1, maximum=2048, step=1, label="Max Tokens", value=256),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature", value=0.7),
        gr.State(value=[])  # Ajout de l'historique
    ],
    outputs=[
        gr.Chatbot(label="History"),  # Utilisation de Chatbot pour l'historique
        gr.State()  # Ajout de l'historique
    ],
    title="Ollama Demo 🐳 🦙 🤗",
    description="""
    Bienvenue sur Docker_Ollama, un espace dédié à l'exploration et au test des modèles Ollama.

    Cette Démo permet aux utilisateurs de tester tous les modèles Ollama dont la taille est inférieure à 10 milliards de paramètres directement depuis cette interface.

    L'Application tourne sur une machine Hugging Face Free Space : 2 CPU - 16Gb RAM
    
    Soyez patient. Chaque Nouveau Modèle Selectionner Devra Etre Télécharger Avant de Répondre.
    
    Il faut compter 2 a 3 minutes pour la réponse d'un modèle 7b alors que quelques dizaines de secondes suffisent pour un modèle 1b.
    
    """
)

def stop_processing():
    global stop_flag
    stop_flag = True

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, pwa=True)