from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

access_token=os.getenv('MODEL_REPO_ID')
model_id = "iGeniusAI/Italia-9B-Instruct-v0.1"

model = AutoModelForCausalLM.from_pretrained(model_id,token=access_token, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

t_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto",
    return_full_text=False, 
    top_p = 0.95, 
    top_k = 50
)

SYSTEM_PROMPT = """Il tuo nome è Modello Italia. Tu sei un'intelligenza artificiale, un modello di linguaggio naturale addestrato da iGenius su Leonardo, uno dei supercomputer più potenti al mondo."""
TEMPERATURE = 0.3
MAX_NEW_TOKENS = 250

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": "Ciao come stai?"},
]

conv_template = tokenizer.apply_chat_template(
        messages,
        tokenize=False
    )

outputs = t_pipeline(
    conv_template,
    max_new_tokens=MAX_NEW_TOKENS,
    do_sample=True,
    temperature=TEMPERATURE,
    num_return_sequences=1,
)
print(outputs[0]["generated_text"])