from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch

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

# Define course information
course_info = {
    "Engineering": ["Civil Engineering", "Mechanical Engineering", "Electrical Engineering", "Software Engineering", "etc."],
    "Information Technology": ["Computer Science", "Information Systems", "Cybersecurity", "Data Science", "etc."],
    "Business": ["Business Administration", "Accounting", "Finance", "Marketing", "Management", "etc."],
    "Health Sciences": ["Nursing", "Pharmacy", "Health Information Management", "Public Health", "etc."],
    "Design and Architecture": ["Architecture", "Industrial Design", "Visual Communication", "etc."],
    "Science": ["Environmental Science", "Biotechnology", "Chemistry", "Physics", "etc."],
    "Law": ["Law", "Legal Studies", "etc."],
    "Arts and Social Sciences": ["Communication", "Education", "International Studies", "Sociology", "etc."],
    "Built Environment": ["Urban Planning", "Construction Management", "Property Economics", "etc."],
    "Creative Industries": ["Animation", "Photography", "Creative Writing", "Film and Television", "etc."]
}

# Function to generate response
def generate_response(input_text, chat_history=[]):
    # Tokenize the new input sentence
    new_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")

    # Append the new user input tokens to the chat history
    bot_input_ids = torch.cat([torch.tensor(chat_history), new_input_ids], dim=-1)

    # Generate a response
    response_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)

    # Decode the generated response
    response = tokenizer.decode(response_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

    # Append the response to the chat history
    chat_history.extend(new_input_ids.tolist())
    chat_history.extend(response_ids[:, bot_input_ids.shape[-1]:].tolist())

    return response, chat_history

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=["text", "text"],
    outputs="text",
    title="AI ChatBot",
    description="A conversational AI powered by DialoGPT.",
    theme="huggingface"
)

# Launch the interface
iface.launch()