Spaces:
Sleeping
Sleeping
File size: 1,990 Bytes
2268915 8507a1e 2268915 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import os
import dotenv
import requests
from transformers import AutoTokenizer
import streamlit as st
dotenv.load_dotenv()
token = os.environ['HF_TOKEN']
# add models here
modelos = {
'Felladrin/Llama-68M-Chat-v1': '[/INST]',
'google/gemma-7b-it': '<start_of_turn>model\n',
'google/gemma-2-2b-it': '<start_of_turn>model\n'
}
nome_modelo = st.selectbox('select a model:', options=modelos)
token_modelo = modelos[nome_modelo]
url = f"https://api-inference.huggingface.co/models/{nome_modelo}"
tokenizer = AutoTokenizer.from_pretrained(nome_modelo, token=token)
if 'modelo_atual' not in st.session_state or st.session_state['modelo_atual'] != nome_modelo:
st.session_state['modelo_atual'] = nome_modelo
st.session_state['mensagens'] = []
mensagens = st.session_state['mensagens']
area_chat = st.empty()
pergunta_usuario = st.chat_input('Do you question here: ')
if pergunta_usuario:
mensagens.append({'role': 'user', 'content': pergunta_usuario})
inputs = tokenizer.apply_chat_template(mensagens, tokenize=False, add_generation_prompt=True)
json = {
'inputs': inputs,
'parameters': {'max_new_tokens': 1_000},
'options': {'use_cache': False, 'wait_for_model': True},
}
headers = {
'Authorization': f'Bearer {token}',
}
# Send request and verify response
response = requests.post(url, json=json, headers=headers).json()
# verify if response is valid
if isinstance(response, list) and len(response) > 0 and 'generated_text' in response[0]:
resposta_chatbot = response[0]['generated_text'].split(token_modelo)[-1]
mensagens.append({'role': 'assistant', 'content': resposta_chatbot})
else:
mensagens.append({'role': 'assistant', 'content': "Desculpe, algo deu errado ao processar sua solicitação."})
with area_chat.container():
for mensagem in mensagens:
chat = st.chat_message(mensagem['role'])
chat.markdown(mensagem['content'])
|