import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Charger le modèle et le tokenizer
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("Hawoly18/llama3.2-3B-Wolof")
    model = AutoModelForCausalLM.from_pretrained("Hawoly18/llama3.2-3B-Wolof")
    
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    return tokenizer, model

tokenizer, model = load_model()

# Fonction de génération de réponse
def generate_response(model, tokenizer, question, max_length=512):
    input_text = f"Question: {question}\nRéponse:"
    input_ids = tokenizer.encode(input_text, return_tensors='pt', padding=True, truncation=True)
    attention_mask = input_ids != tokenizer.pad_token_id

    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=max_length,
            attention_mask=attention_mask,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            num_beams=5,  # Recherche par faisceaux pour améliorer la qualité
            no_repeat_ngram_size=2,  # Éviter la répétition des n-grammes
            early_stopping=True
        )
    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    response = response.replace(input_text, "").strip()  # Extraire la réponse
    return response

# Interface Streamlit
st.title("Modèle de génération de réponses")
st.write("Entrez une question pour obtenir une réponse générée par le modèle.")

# Input de l'utilisateur
question = st.text_input("Votre question:")

if question:
    with st.spinner("Génération de la réponse..."):
        response = generate_response(model, tokenizer, question)
        st.write("Réponse:")
        st.write(response)

# Option pour partager l'interface sur Streamlit Cloud ou en local