import gradio as gr
import spaces 
## Load model directly
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import HfApi

# Récupérer le jeton d'accès API depuis Hugging Face Secrets
api = HfApi()
secret = api.secrets.get("alex-abb/8B-on-GPU-ZERO/API_KEY")

# Utiliser le jeton d'accès API pour charger le modèle
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", token=secret["value"])
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", token=secret["value"])





@spaces.GPU(duration=120)

# Fonction de génération de texte
def generate_text(prompt):
    
    inputs = tokenizer(prompt, return_tensors="pt")
    
    response_ids = model.generate(inputs.input_ids)
    
    response_text = tokenizer.decode(response_ids[0], skip_special_tokens=True)

    return response_text 
# Définir une fonction pour l'interface de chat
def chatbot(message,history):
    return generate_text(message)


gr.ChatInterface(chatbot).launch()