import torch
import pandas as pd
import streamlit as st
from transformers import pipeline
MODEL_NAME = 'Qwen/Qwen2.5-1.5B-Instruct'
df = pd.read_csv('hf://datasets/lhoestq/pokemonData/data/Pokemon.csv')
df.drop_duplicates(subset='Name')
st.write('Warning: Incredibly slow application, but it works. Woo AI.')
pokemon = st.selectbox("Choose a Pokémon", df['Name'], index=24)
chat = [
    {"role": "user",
     "content": f"You are a friendly chatbot who always responds in the style of {pokemon} from Pokémon."},
]
pipe = pipeline("text-generation", model=MODEL_NAME)
pipe(chat, max_new_tokens=128)
if pokemon:
    prompt = st.text_area(f'Say something to {pokemon}:') 
    if prompt:
        chat.append(
            {
                'role': 'user',
                'content': prompt
            }
        )
        out = pipe(chat, max_new_tokens=128)
        st.subheader(f'{pokemon} says:')
        poke_out = st.write(out[0]['generated_text'][-1]['content'])
        st.write(poke_out)