File size: 1,414 Bytes
70f9aa7 |
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 |
import os
from dotenv import load_dotenv
import streamlit as st
from transformers import AutoTokenizer
import requests
load_dotenv()
modelos = {
"mistralai/Mixtral-8x7B-Instruct-v0.1": "[/INST]",
"google/gemma-7b-it": "<start_of_turn>model\n",
}
modelo = st.selectbox("Selecion um modelo: ", options=modelos)
token_modelo = modelos[modelo]
if "modelo_atual" not in st.session_state or st.session_state["modelo_atual"] != modelo:
st.session_state["modelo_atual"] = modelo
st.session_state["mensagens"] = []
nome_modelo = st.session_state["modelo_atual"]
tokenizer = AutoTokenizer.from_pretrained(nome_modelo, token=os.getenv("HF_KEY"))
url = f"https://api-inference.huggingface.co/models/{nome_modelo}"
headers = {"Authorization": f"Bearer {os.getenv('HF_KEY')}"}
mensagens = st.session_state["mensagens"]
area_chat = st.empty()
pergunta_usuario = st.chat_input("Faça a sua pergunta aqui:")
if pergunta_usuario:
mensagens.append({"role": "user", "content": pergunta_usuario})
template = tokenizer.apply_chat_template(
mensagens, tokenize=False, add_generation_prompt=True
)
json = {
"inputs": template,
"parameters": {"max_new_tokens": 1000},
"options": {"use_cache": False, "wait_for_model": True},
}
res = requests.post(url, json=json, headers=headers).json()
mensagem_chat = res[0]["generated_text"].split(token_modelo)[-1]
|