Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoTokenizer, AutoModel, GPT2LMHeadModel, GPT2Tokenizer | |
import torch | |
from sklearn.metrics.pairwise import cosine_similarity | |
import numpy as np | |
# Load the bi-encoder model and tokenizer | |
bi_encoder_model_name = "nasa-impact/nasa-smd-ibm-st-v2" | |
bi_tokenizer = AutoTokenizer.from_pretrained(bi_encoder_model_name) | |
bi_model = AutoModel.from_pretrained(bi_encoder_model_name) | |
# Load the GPT-2 model and tokenizer for response generation | |
gpt2_model_name = "gpt2" | |
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name) | |
gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name) | |
def encode_text(text): | |
inputs = bi_tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128) | |
outputs = bi_model(**inputs) | |
# Ensure the output is 2D by averaging the last hidden state along the sequence dimension | |
return outputs.last_hidden_state.mean(dim=1).detach().numpy() | |
def generate_response(user_input): | |
# Encode the user input | |
user_embedding = encode_text(user_input) | |
# Generate a response using GPT-2 | |
gpt2_inputs = gpt2_tokenizer.encode(user_input, return_tensors='pt') | |
gpt2_outputs = gpt2_model.generate(gpt2_inputs, max_length=150, num_return_sequences=1) | |
generated_text = gpt2_tokenizer.decode(gpt2_outputs[0], skip_special_tokens=True) | |
return generated_text | |
def chatbot(user_input): | |
response = generate_response(user_input) | |
return response | |
# Create the Gradio interface | |
iface = gr.Interface( | |
fn=chatbot, | |
inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."), | |
outputs="text", | |
title="Dynamic Response Chatbot", | |
description="A chatbot using a bi-encoder model to understand the input and GPT-2 to generate dynamic responses." | |
) | |
# Launch the interface | |
iface.launch() | |