Testing_BrokenSpace / app.py_Bloom1B_02_OK
MisterAI's picture
Rename app.py to app.py_Bloom1B_02_OK
c7eb30c verified
#testing bloom1b 02
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,
max_length=100,
num_return_sequences=1,
no_repeat_ngram_size=2,
do_sample=True,
top_p=0.9,
temperature=0.7
)
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()