from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import streamlit as st
import pandas as pd
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained("Celestinian/TopicGPT")
tokenizer = AutoTokenizer.from_pretrained("Celestinian/TopicGPT")


def get_topic(prompt, temperature=0.01, max_size=20):
    input_ids = tokenizer.encode("#CONTEXT# " + prompt + " #TOPIC#", return_tensors='pt')
    input_ids = input_ids.to(device)
    model.eval()
    model.to(device)

    output_tokens = []
    eos_token_id = tokenizer.encode('#')[0]
    print(input_ids)
    for _ in range(max_size):
        with torch.no_grad():
            outputs = model(input_ids)
        logits = outputs.logits[:, -1, :] / temperature
        next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
        if next_token.item() == eos_token_id:
            break
        input_ids = torch.cat((input_ids, next_token), dim=-1)
        output_tokens.append(next_token.item())

    print(output_tokens)
    output = tokenizer.decode(output_tokens)
    clean_output = output.replace('\n', '')
    print(output)
    return clean_output


text = st.text_area("Review", "Trustpilot is the one and only you can trust and share your thoughts with others.Love it !")
outputs = get_topic(text)
st.write(outputs)