import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

base_model_id = "elvyslpontes/finnlp-finllm-task-1-classification"
model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(
    base_model_id,
    model_max_length=512,
    padding_side="left",
    add_eos_token=True,
)

query = "Analyze sentences from earnings conference calls and identify their argumentative function. Each sentence is either a premise, offering evidence or reasoning, or a claim, asserting a conclusion or viewpoint. Return only premise or claim. Text: I mean, sometimes it's not that you came up with some brilliant strategy, it's just like really good work consistently over a long period of time. Answer:"

runtimeFlag = "cuda:0"
inputs = tokenizer(query, return_tensors="pt").to(runtimeFlag)
generated_ids = model.generate(
    **inputs,
    max_new_tokens=3,
    pad_token_id=tokenizer.eos_token_id,
)
decoded = tokenizer.batch_decode(generated_ids)
answer = decoded[0].split("Answer:")[1].split("</s>")[1].lower()

print(answer)