from transformers import pipeline
import gradio as gr  # Import Gradio for UI

# Load a text-generation model
chatbot = pipeline("text-generation", model="microsoft/DialoGPT-medium")

# Load the classification model
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Customize the bot's knowledge base with predefined responses
faq_responses = {
    "study tips": "Here are some study tips: 1) Break your study sessions into 25-minute chunks (Pomodoro Technique). 2) Test yourself frequently. 3) Stay organized using planners or apps like Notion or Todoist.",
    "resources for studying": "You can find free study resources on websites like Khan Academy, Coursera, and edX. For research papers, check Google Scholar.",
    "how to focus": "To improve focus, try studying in a quiet place, remove distractions like your phone, and use apps like Forest or Focus@Will.",
    "time management tips": "Start by creating a to-do list each morning. Prioritize tasks using methods like Eisenhower Matrix and allocate specific time blocks for each task.",
    "how to avoid procrastination": "Break tasks into smaller steps, set deadlines, and reward yourself after completing milestones. Tools like Trello can help you stay organized."
}

# Define the chatbot's response function
def faq_chatbot(user_input):
    # Classify user input based on predefined FAQ categories
    classified_user_input = classifier(user_input, candidate_labels=list(faq_responses.keys()))

    # Get the highest confidence score label
    predicted_label = classified_user_input["labels"][0]
    confidence_score = classified_user_input["scores"][0]

    # Confidence threshold (adjust as needed)
    threshold = 0.5

    # If classification confidence is high, return the corresponding FAQ response
    if confidence_score > threshold:
        return faq_responses[predicted_label]

    # If no FAQ match, use the AI model to generate a response
    conversation = chatbot(user_input, max_length=50, num_return_sequences=1)
    return conversation[0]['generated_text']

# Create the Gradio interface
interface = gr.Interface(
    fn=faq_chatbot,  # Function to process user input
    inputs=gr.Textbox(lines=2, placeholder="Ask me about study tips, resources, or time management..."),  # Input field
    outputs="text",  # Output text
    title="Student FAQ Chatbot",
    description="Ask me study tips, time management strategies, or where to find good study resources!"
)

# Launch the chatbot and make it accessible via a public Gradio link
interface.launch(share=True)