Spaces:
Running
Running
import gradio as gr | |
from transformers import AutoTokenizer, TFBlenderbotForConditionalGeneration | |
import tensorflow as tf | |
print("Loading the model......") | |
model_name = "WICKED4950/Irisonego5" | |
strategy = tf.distribute.MirroredStrategy() | |
tf.config.optimizer.set_jit(True) # Enable XLA | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
with strategy.scope(): | |
model = TFBlenderbotForConditionalGeneration.from_pretrained(model_name) | |
print("Interface getting done....") | |
# Define the chatbot function | |
def predict(user_input): | |
# Tokenize input text | |
inputs = tokenizer(user_input, return_tensors="tf", padding=True, truncation=True) | |
# Generate the response using the model | |
response_id = model.generate( | |
inputs['input_ids'], | |
max_length=128, # Set max length of response | |
do_sample=True, # Sampling for variability | |
top_k=15, # Consider top 50 tokens | |
top_p=0.95, # Nucleus sampling | |
temperature=0.8 # Adjusts creativity of response | |
) | |
# Decode the response | |
response = tokenizer.decode(response_id[0], skip_special_tokens=True) | |
return response | |
# Gradio interface | |
gr.Interface( | |
fn=predict, | |
inputs=gr.Textbox(label="Ask Iris anything!"), | |
outputs=gr.Textbox(label="Iris's Response"), | |
examples=[ | |
["What should I do if I'm feeling down?"], | |
["How do I deal with stress?"], | |
["Tell me something positive!"] | |
], | |
description="A chatbot trained to provide friendly and comforting responses. Type your question below and let Iris help!", | |
title="Iris - Your Friendly Mental Health Assistant", | |
).launch() | |