Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
import torch | |
# Load the model and tokenizer | |
model_name = 'FridayMaster/fine_tune_embedding' # Replace with your model's repository name | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForSequenceClassification.from_pretrained(model_name) # Use the appropriate class | |
# Define a function to generate responses | |
def generate_response(prompt): | |
# Tokenize the input prompt | |
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512) | |
with torch.no_grad(): | |
# Get the model output | |
outputs = model(**inputs) | |
# Process the output logits | |
logits = outputs.logits | |
predicted_class_id = logits.argmax().item() | |
# Generate a response based on the predicted class | |
response = f"Predicted class ID: {predicted_class_id}" | |
return response | |
# Create a Gradio interface | |
iface = gr.Interface( | |
fn=generate_response, | |
inputs=gr.Textbox(label="Enter your message", placeholder="Type something here..."), | |
outputs=gr.Textbox(label="Response"), | |
title="Chatbot Interface", | |
description="Interact with the fine-tuned chatbot model." | |
) | |
# Launch the Gradio app | |
if __name__ == "__main__": | |
iface.launch() | |