Update README.md
#1
by
SafeerChalil
- opened
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)