app.py

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

# Define a dictionary of model names and their corresponding Hugging Face model IDs
models = {
    "GPT-Neo-125M": "EleutherAI/gpt-neo-125M",
    "GPT-J-6B": "EleutherAI/gpt-j-6B",
    "GPT-NeoX-20B": "EleutherAI/gpt-neox-20b",
    "GPT-3.5-Turbo": "gpt2",  # Placeholder for illustrative purposes
}

# Initialize tokenizers and models
tokenizers = {}
models_loaded = {}

for model_name, model_id in models.items():
    tokenizers[model_name] = AutoTokenizer.from_pretrained(model_id)
    models_loaded[model_name] = AutoModelForCausalLM.from_pretrained(model_id)

def chat(model_name, user_input, history=[]):
    tokenizer = tokenizers[model_name]
    model = models_loaded[model_name]
    
    # Encode the input
    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
    
    # Generate a response
    with torch.no_grad():
        output = model.generate(input_ids, max_length=150, pad_token_id=tokenizer.eos_token_id)
    
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    # Clean up the response to remove the user input part
    response = response[len(user_input):].strip()
    
    # Append to chat history
    history.append((user_input, response))
    
    return history, history

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Chat with Different Models")
    
    model_choice = gr.Dropdown(list(models.keys()), label="Choose a Model")
    chatbot = gr.Chatbot(label="Chat")
    message = gr.Textbox(label="Message")
    submit = gr.Button("Submit")
    
    submit.click(chat, inputs=[model_choice, message, chatbot], outputs=[chatbot, chatbot])
    message.submit(chat, inputs=[model_choice, message, chatbot], outputs=[chatbot, chatbot])

# Launch the demo
demo.launch()