import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Configuração de dispositivo
DEVICE = 0 if torch.cuda.is_available() else -1
TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32

# Carrega modelo A - LLaMA 2
model_a = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    torch_dtype=TORCH_DTYPE
)
tokenizer_a = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

pipe_a = pipeline(
    "text-generation",
    model=model_a,
    tokenizer=tokenizer_a,
    device=DEVICE,
    return_full_text=False,
    pad_token_id=tokenizer_a.eos_token_id
)

# Carrega modelo B - Falcon 7B Instruct (sem autenticação)
model_b = AutoModelForCausalLM.from_pretrained(
    "tiiuae/falcon-7b-instruct",
    torch_dtype=TORCH_DTYPE
)
tokenizer_b = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")

pipe_b = pipeline(
    "text-generation",
    model=model_b,
    tokenizer=tokenizer_b,
    device=DEVICE,
    return_full_text=False,
    pad_token_id=tokenizer_b.eos_token_id
)

# Funções auxiliares para formatar o prompt
def format_llama_prompt(user_input):
    return f"[INST] <<SYS>>\nVocê é um assistente útil.\n<</SYS>>\n\n{user_input.strip()} [/INST]"

def format_falcon_prompt(user_input):
    return f"Responda em português: {user_input.strip()}"

# Interface simples para testar os modelos
if __name__ == "__main__":
    while True:
        prompt = input("\nDigite uma pergunta (ou 'sair'): ").strip()
        if prompt.lower() == "sair":
            break

        print("\n=== Resposta do LLaMA 2 ===")
        llama_response = pipe_a(format_llama_prompt(prompt), max_new_tokens=200)[0]['generated_text']
        print(llama_response)

        print("\n=== Resposta do Falcon 7B ===")
        falcon_response = pipe_b(format_falcon_prompt(prompt), max_new_tokens=200)[0]['generated_text']
        print(falcon_response)