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

Load AI model and tokenizer

model_name = "gpt2" # Using GPT-2 (lightweight)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

Conversation history storage

conversation_history = []

Function to generate AI responses while remembering past messages

def chatbot_response(user_input):
global conversation_history # Keep track of previous messages

# Append new user input to history
conversation_history.append(f"User: {user_input}")

# Keep only the last 5 exchanges to avoid memory issues
if len(conversation_history) > 5:
    conversation_history.pop(0)

# Format conversation as context
context = " ".join(conversation_history)

# Generate AI response
input_ids = tokenizer(context, return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_length=200)
response = tokenizer.decode(output[0], skip_special_tokens=True)

# Save AI response to history
conversation_history.append(f"AI: {response}")

return response

Create Gradio chatbot UI

def chat_interface(user_input):
return chatbot_response(user_input)

gr.Interface(fn=chat_interface, inputs="text", outputs="text", title="Memory Chatbot").launch(share=True)

Ready to merge
This branch is ready to get merged automatically.
Your need to confirm your account before you can post a new comment.

Sign up or log in to comment