app.py

import gradio as gr
import spaces
from huggingface_hub import login
import accelerate
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
import os
import torch
from typing import Optional, Iterator, Dict, Any, List
from threading import Thread
from types import NoneType
import traceback


# Initialize logging and device information
print(f"Is CUDA available: {torch.cuda.is_available()}")
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")

MAX_NEW_TOKENS = 2**13
DEFAULT_MAX_NEW_TOKENS = 0.65*MAX_NEW_TOKENS
DEFAULT_SYSTEM_PROMPT = """
Tu es un expert en extraction de données dans des documents très longs et bruités. 
Tu comprends le sujet grâce à des liens sémantiques que tu peux extraire. 
Tu sers à créer des concepts hiérarchiques ainsi que des liens entre ceux-ci.
Réponds de manière claire et formelle et va droit au but dans ta tâche.
"""

class HuggingFaceLogin:
    """Handles authentication to the Hugging Face Hub using environment variables or explicit tokens."""
    def __init__(self, env_token_key: str = "HF_TOKEN"):
        """Initialize the login handler.

        Args:
            env_token_key (str): Environment variable key containing the token. Defaults to "HF_TOKEN".
        """
        self.token = os.getenv(env_token_key)
        
    def login(self, token: str = None) -> bool:
        """Authenticate with the Hugging Face Hub.

        Args:
            token (Optional[str]): Optional explicit token. If not provided, uses token from environment.

        Returns:
            bool: True if login successful, False otherwise.
        
        Raises:
            ValueError: If no token is available (neither in env nor passed explicitly).
        """

        if not self.token:
            raise ValueError("No authentication token provided. Set HF_TOKEN environment variable or pass token explicitly.")
        try:
            print("Logging in to the Hugging Face Hub...")
            login(token=self.token)
            return True
        except Exception as e:
            print(f"Login failed: {str(e)}")
            return False

model_config_4bit = BitsAndBytesConfig(
    load_in_4bit                = True,
    bnb_4bit_use_double_quant   = True,
    bnb_4bit_quant_type         = "nf4",
    bnb_4bit_compute_dtype=torch.float16
)

model_config_8bit = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
    llm_int8_has_fp16_weight=False,
    bnb_8bit_compute_dtype=torch.float16
)

if torch.cuda.is_available():
    model_id = "meta-llama/Llama-3.1-8B-Instruct"
    model = AutoModelForCausalLM.from_pretrained(model_id,
                                                 quantization_config=model_config_8bit,
                                                 device_map="auto")
    tokenizer = AutoTokenizer.from_pretrained(model_id)

# Helper function to generate responses from the LLM
def generate_llm_response(
    conversation: List[Dict[str, str]],
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    repetition_penalty: float
) -> str:
    """Generate a response from the LLM based on the conversation."""
    input_ids = tokenizer.apply_chat_template(
        conversation, 
        return_tensors="pt",
        add_generation_prompt=True
    )
    
    input_ids = input_ids.to(model.device)
    
    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=2*60.0,
        skip_prompt=True,
        skip_special_tokens=True
    )
    
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
        pad_token_id=tokenizer.eos_token_id,
    )
    
    t = Thread(
        target=model.generate,
        kwargs=generate_kwargs
    )
    t.start()
    
    # Collect the output
    accumulated_response = ""
    for text in streamer:
        accumulated_response += text
        yield accumulated_response

def append_text_knowledge(file_path: str) -> str:
    """
    Reads content from a selected file and returns it as a string.
    
    Args:
        file_path (str): Path to the selected file
        
    Returns:
        str: Content of the file or empty string if no file selected
    """
    if file_path:
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()
        except Exception as e:
            print("Error reading file: {e}")
            return ""
    return ""

knowledge_textbox = gr.Textbox(
    label="Knowledge Text",
    lines= 20,
    visible=False
)

with gr.Blocks() as demo:
    gr.Markdown("# Ontology Generation with Chain-of-Thought")
    chatbot = gr.Chatbot(type="messages")
    message_input  = gr.Textbox(
        label="message",
        placeholder="Ask about the elicitation text...",
        lines=2,
        submit_btn=True
    )
    
    with gr.Row():
        file_explorer = gr.FileExplorer(
            glob="**/*.txt",
            file_count="single",
            label="Upload file",
            show_label=True
        )
        knowledge_input = gr.Textbox(
            label="Knowledge text",
            lines=6,
            visible=True
        )
    with gr.Accordion("Advanced Settings", open=False):
        system_prompt_input = gr.Textbox(
            label="System Prompt",
            lines=4,
            value=DEFAULT_SYSTEM_PROMPT
        )
        with gr.Row():
            with gr.Column():
                max_tokens_slider = gr.Slider(
                    label="Max new tokens",
                    minimum=1,
                    maximum=MAX_NEW_TOKENS,
                    step=1,
                    value=DEFAULT_MAX_NEW_TOKENS
                )
                temperature_slider = gr.Slider(
                    label="Temperature",
                    minimum=0.1,
                    maximum=4.0,
                    step=0.1,
                    value=0.2
                )
            
            with gr.Column():
                top_p_slider = gr.Slider(
                    label="Top-p (nucleus sampling)",
                    minimum=0.05,
                    maximum=1.0,
                    step=0.05,
                    value=0.8
                )
                top_k_slider = gr.Slider(
                    label="Top-k",
                    minimum=1,
                    maximum=1000,
                    step=1,
                    value=50
                )
                repetition_penalty_slider = gr.Slider(
                    label="Repetition penalty",
                    minimum=1.0,
                    maximum=2.0,
                    step=0.05,
                    value=1.0
                )

    # Example prompts
    examples = gr.Examples(
        examples=[
            ["Extract meaningful entities in your knowledge document in order to create a Turtle-formatted output where you create classes and sub-classes and object properties automatically."],
            ["Make a simple list of the classes, sub-classes and object properties that can be extracted from the knowledge document."]
        ],
        inputs=message_input
    )
    
    def user_message(message:str, history:List[Dict[str, str]]):
        """Add user message to chat history.

        Args:
            message (str): The User Message to send
            history (List[Dict[str,str]]): The previous chat conversation history.
        """
        if message.strip() == "":
            return history, message
    
        history = history + [{"role":"user", "content": message}]
        return history, ""

    def bot_response(history, knowledge, system_prompt, max_tokens, temp, top_p, top_k, rep_penalty):
        """Generate assistant response with visible thinking.

        Args:
            history (List[Dict[str, str]]): The previous chat conversation history
            knowledge (Any): Documents to pass as knowledge to the multimodal model
            system_prompt (str): System prompt that the model follows
            max_tokens (int): Max number of allowed output tokens
            temp (float): Model's Temperature
            top_p (int): Model's Top p value
            top_k (int): Model's Top k value
            rep_penalty (float): Model's repetition penalty
            
        Returns:
            history (List[Dict[str, str]]): The history of the conversation updated
        """
        try:
            if not history or history[-1]["role"] != "user":
                return history
            
            user_message = history[-1]["content"]
            # thinking message with pending status
            history.append({
                "role": "assistant",
                "content": "Je réfléchis étape par étape...",
                "metadata": {
                    "title": "Réflexion",
                    "status": "pending"
                }
            })
            yield history
            
            thinking_conversation = []
            if system_prompt:
                thinking_conversation.append({"role": "system", "content": system_prompt})
            if knowledge:
                thinking_conversation.append({
                    "role": "assistant",
                    "content": f"Voici le document que je dois comprendre: {knowledge}\n\nJe vais l'analyser étape par étape."
                })
                
            for msg in history[:-2]: # All msg except user message and thinking part
                thinking_conversation.append(msg)
            
            thinking_prompt = user_message + "\n\nRéfléchis étape par étape. D'abord identifie l'intention de l'utilisateur. Quand tu as compris ce qui t'est demandé, commence à établir un plan clair et précis que tu peux suivre. Utilise l'italic et le gras en Markdown pour séquencer et prioriser tes actions."
            thinking_conversation.append({"role": "user", "content": thinking_prompt})
            
            # GENERATE THINKING
            for thinking_partial in generate_llm_response(thinking_conversation,
                                                          max_new_tokens=max_tokens * 2,
                                                          temperature=temp,
                                                          top_p=top_p,
                                                          top_k=top_k,
                                                          repetition_penalty=rep_penalty):
                # update the thinking message
                history[-1] = {
                    "role": "assistant",
                    "content": thinking_partial,
                    "metadata": {
                        "title": "Réflexion",
                        "status": "done"
                    }
                }
                yield history
            
            history[-1]["metadata"]["status"] = "done"
            yield history
            
            print("DEBUG:\t\tYielded history of ```thinking_result```")

            final_conversation = []
            if system_prompt:
                    final_conversation.append({"role": "system", "content": system_prompt})
            if knowledge:
                final_conversation.append({
                    "role": "assistant",
                    "content": f"J'ai analysé ce document: {knowledge}"
                })
            
            for msg in history[:-1]: # exclude thinking
                if "metadata" not in msg or "title" not in msg.get("metadata", {}):
                    final_conversation.append(msg)

            final_conversation.append({
                "role": "assistant",
                "content": f"Voici mon analyse étape par étape:\n{history[-1]['content']}\n\nMaintenant je vais formaliser le résultat final."
            })
            final_conversation.append({
                "role": "assistant",
                "content": "Je formule ma réponse finale..."
            })
            yield history
            
            for final_partial in generate_llm_response(final_conversation,
                                                       max_new_tokens=max_tokens,
                                                       temperature=temp * 0.8,  # Lower temperature for final answer
                                                       top_p=top_p,
                                                       top_k=top_k,
                                                       repetition_penalty=rep_penalty):
                history[-1]["content"] = final_partial
                yield history
            print("DEBUG:\t\tYielded history of ```final_answer```")

        except Exception as e:
            error_traceback = traceback.format_exc()
            print(f"Error traceback:\n{error_traceback}")
            
            history.append({
                "role": "assistant",
                "content": f"An error occurred: {str(e)}\n\nTraceback details:\n{error_traceback}"
            })
            yield history
    
    file_explorer.change(
        append_text_knowledge,
        file_explorer,
        knowledge_input
    )
    
    message_input.submit(
        user_message,
        inputs=[message_input, chatbot],
        outputs=[chatbot, message_input]
    ).then(
        bot_response,
        inputs=[
            chatbot,
            knowledge_input,
            system_prompt_input,
            max_tokens_slider,
            temperature_slider,
            top_p_slider,
            top_k_slider,
            repetition_penalty_slider
        ],
        outputs=chatbot
    )
if __name__ == "__main__":
    auth = HuggingFaceLogin()
    if auth.login():
        print("Login successful!")
    demo.queue().launch()