import json
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# Load the JSON data from the file
with open('uts_courses.json') as f:
    data = json.load(f)

# Load the question-answering pipeline
qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")

# Load tokenizer and model for dialog generation
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")

# Main function to interact with the user
def main():
    print("Welcome! I'm the UTS Course Chatbot. How can I assist you today?")
    print("You can ask questions about UTS courses or type 'exit' to end the conversation.")

    while True:
        user_input = input("You: ")
        if user_input.lower() == 'exit':
            print("Bot: Exiting the program.")
            break

        if "courses" in user_input.lower() and "available" in user_input.lower():
            field = user_input.split("in ")[-1]
            courses = data['courses'].get(field, [])
            if courses:
                response = f"Courses in {field}: {', '.join(courses)}"
            else:
                response = f"No courses found in {field}."
        else:
            response = generate_dialog_response(user_input)
        print("Bot:", response)

# Function to generate response using dialog generation model
def generate_dialog_response(user_input):
    input_text = user_input + tokenizer.eos_token
    input_ids = tokenizer.encode(input_text, return_tensors="pt")

    # Generate response
    response_ids = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response_text = tokenizer.decode(response_ids[0], skip_special_tokens=True)

    return response_text

if __name__ == "__main__":
    main()