aquibmoin's picture
Update app.py
11d423c verified
raw
history blame
1.84 kB
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()