Chatbot2 / chatbot.py
jithenderchoudary's picture
Create chatbot.py
a73d81f verified
from transformers import BertTokenizer, BertForSequenceClassification
import torch
from simple_salesforce import Salesforce
from model import label_encoder # Assuming the label_encoder was saved in the model.py file
# Load pre-trained model and tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("results") # Path to your trained model
# Salesforce authentication setup
sf = Salesforce(username="your_username", password="your_password", security_token="your_token")
# Function to classify the intent of user input
def classify_intent(user_input):
# Tokenize the user input and make a prediction
encoded_input = tokenizer(user_input, padding=True, truncation=True, return_tensors="pt")
# Run the model without gradients for efficiency
with torch.no_grad():
output = model(**encoded_input)
# Get the predicted label and return it
predicted_label = torch.argmax(output.logits, dim=1).item()
return label_encoder.inverse_transform([predicted_label])[0] # Return intent label
# Function to generate a response based on the classified intent
def generate_response(intent):
# Simple responses based on predefined intents
responses = {
"greeting": "Hi, how can I help you today?",
"goodbye": "Goodbye! Have a nice day!",
"question": "Hugging Face is a company that specializes in NLP models."
}
# Return response for the given intent
return responses.get(intent, "Sorry, I didn't understand that.")
# Function to create a case in Salesforce based on user input and chatbot response
def create_case_in_salesforce(user_input, response):
case_data = {
'Subject': f"User Query: {user_input}",
'Description': f"User asked: {user_input}\nResponse: {response}"
}
# Create the case in Salesforce
case = sf.Case.create(case_data)
print(f"Case Created with ID: {case['id']}")
# Main function to handle the full interaction
def chatbot_response(user_input):
# Step 1: Classify the user input into an intent
intent = classify_intent(user_input)
# Step 2: Generate the appropriate response based on the intent
response = generate_response(intent)
# Step 3: Optionally create a case in Salesforce (e.g., for queries or support requests)
create_case_in_salesforce(user_input, response)
# Step 4: Return the chatbot's response
return response
# Example of using the chatbot
if __name__ == "__main__":
user_input = "Hello, tell me about Hugging Face."
print(chatbot_response(user_input)) # Should print the chatbot's response and create a Salesforce case