CTP-Week-9 / app.py
alyliann's picture
Update app.py
5d3f44a verified
raw
history blame contribute delete
992 Bytes
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)