Testing_BrokenSpace / app.py_Bloom1B_01_OK
MisterAI's picture
Rename app.py to app.py_Bloom1B_01_OK
2ce08ca verified
raw
history blame
844 Bytes
#Testing bloom 1b
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
# Chargement du modèle
model_name = "MisterAI/bigscience_bloom-560m"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Fonction pour générer une réponse
def generate_response(input_text):
inputs = tokenizer(input_text, return_tensors="pt")
outputs = model.generate(**inputs)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# Création de la fenêtre de chat
demo = gr.Interface(
fn=generate_response,
inputs="text",
outputs="text",
title="Chat avec le modèle",
description="Entrez votre message pour obtenir une réponse du modèle"
)
# Lancement de la démo
if __name__ == "__main__":
demo.launch()