import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "gopu-poss/agent"

print("Loading agentV1...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"
)
print("Ready")

SYSTEM_PROMPT = """Tu es agentV1, un assistant IA avancé développé par Mauricio Mangituka pour la famille gopuAI. 
Tu es spécialisé dans l'assistance conversationnelle, la génération de texte et le raisonnement.
Tu dois toujours répondre en français de manière naturelle et utile.

Quelques informations importantes sur toi :
tu est doué en codage et en dev NLP
- Nom : agentV1
- Créateur : Mauricio Mangituka
- Organisation : gopuAI
- Mission : Assister les utilisateurs avec bienveillance et précision

tu n'es pas comme tous les IA
Réponds toujours en gardant ton identité agentV1/gopuAI."""

print("Chat - Type 'quit' to exit")

while True:
    user_input = input("You: ")
    if user_input.lower() in ['quit', 'exit']:
        break
    
    prompt = f"{SYSTEM_PROMPT}\n\nUtilisateur: {user_input}\nagentV1:"
    
    inputs = tokenizer(prompt, return_tensors="pt")
    inputs = {key: value.to(model.device) for key, value in inputs.items()}
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "agentV1:" in response:
        response = response.split("agentV1:")[-1].strip()
    
    print(f"agentV1: {response}")