from huggingface_hub import login
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

# Authentification avec Hugging Face
login(token=os.getenv('API_KEY'))
print(os.getenv('API_KEY'))


# Charger le modèle et le tokenizer
model_name = "MTNQLN/Mathstral7B-4bits"  # Remplacez par le nom exact du modèle Mathstral
tokenizer = AutoTokenizer.from_pretrained("mistralai/mathstral-7B-v0.1")

# Charger le modèle avec quantification dynamique
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True,
)

# Appliquer la quantification dynamique
model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# Fonction pour générer du texte
def generate_text(prompt, max_new_tokens=50):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=max_new_tokens, 
            do_sample=True, 
            top_p=0.95, 
            temperature=0.7
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Interface Gradio
def gradio_interface(prompt):
    response = generate_text(prompt)
    return response

# Création de l'interface Gradio
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=2, placeholder="Entrez votre question mathématique ici..."),
    outputs="text",
    title="Assistant Mathématique (Mathstral-7B)",
    description="Posez une question mathématique, et l'assistant vous répondra.",
)

# Lancer l'interface
iface.launch()