import gradio as gr
import torch
from transformers import (
AutoTokenizer,
AutoModelForCausalLM,
pipeline
)
import spaces
import time
import os
# Model configuration
MODEL_NAME = "microsoft/DialoGPT-medium" # 1.5B parameters, close to 2B
# Alternative 2B models you could try:
# "microsoft/Phi-2" (2.7B - requires special handling)
# "EleutherAI/gpt-neo-2.7B" (2.7B parameters)
# Global variables
tokenizer = None
model = None
chat_history = []
def load_model():
"""Load the model and tokenizer"""
global tokenizer, model
if tokenizer is None or model is None:
print("Loading model and tokenizer...")
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="left")
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Load model
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float32, # Use float32 for CPU compatibility
low_cpu_mem_usage=True
)
print(f"Model {MODEL_NAME} loaded successfully!")
return tokenizer, model
@spaces.GPU(duration=120) # Use GPU if available, with 2-minute timeout
def generate_response(user_message, history=None):
"""
Generate response using the loaded model
Args:
user_message (str): User's input message
history (list): Previous chat history
Returns:
str: Generated response
"""
if history is None:
history = []
try:
# Load model if not already loaded
load_model()
# Prepare input
chat_history = history.copy()
chat_history.append(user_message)
# Combine all messages for context
context = "\n".join([f"Human: {msg}" if i % 2 == 0 else f"Assistant: {msg}"
for i, msg in enumerate(chat_history)])
context += "\nAssistant:"
# Tokenize input
inputs = tokenizer.encode(context, return_tensors="pt", max_length=1024, truncation=True)
# Generate response
with torch.no_grad():
outputs = model.generate(
inputs,
max_length=inputs.shape[1] + 100,
num_return_sequences=1,
temperature=0.7,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.encode("Human")[0]
)
# Decode response
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract only the assistant's response
if "Assistant:" in response:
response = response.split("Assistant:")[-1].strip()
else:
# Fallback: try to extract meaningful response
response = response[len(context):].strip()
# Clean up response
response = response.split("\n")[0].strip()
# Remove any remaining "Human:" parts
if "Human:" in response:
response = response.split("Human:")[0].strip()
# Ensure response is not empty
if not response or len(response.strip()) < 2:
response = "I'm here to chat! What would you like to talk about?"
return response
except Exception as e:
print(f"Error generating response: {e}")
return "I apologize, but I'm having trouble generating a response right now. Please try again!"
def chat_interface(message, history):
"""
Chat interface function
"""
if not message.strip():
return history, ""
# Generate response
response = generate_response(message, history)
# Update history
history.append(message)
history.append(response)
# Keep history manageable (last 10 exchanges)
if len(history) > 20:
history = history[-20:]
return history, ""
def clear_chat():
"""Clear the chat history"""
return []
# Create the Gradio interface
def create_demo():
"""Create the Gradio demo"""
# Custom CSS for better styling
css = """
.gradio-container {
max-width: 800px !important;
margin: auto !important;
}
.header {
text-align: center;
padding: 20px;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
border-radius: 10px;
margin-bottom: 20px;
}
.model-info {
text-align: center;
padding: 10px;
background-color: #f0f2f6;
border-radius: 5px;
margin-bottom: 20px;
font-size: 0.9em;
}
"""
with gr.Blocks(css=css, title="Free 2B Parameter Chatbot") as demo:
# Header
gr.HTML("""
""")
# Model info
gr.HTML(f"""
Model: {MODEL_NAME} (1.5B parameters)
Type: Conversational AI
Powered by: Hugging Face Transformers
""")
# Chat interface
chatbot = gr.Chatbot(
label="Chat with AI",
height=600,
bubble_full_width=False,
avatar_images=(None, None)
)
msg = gr.Textbox(
label="Your message",
placeholder="Type your message here...",
scale=4
)
with gr.Row():
send_btn = gr.Button("Send", variant="primary", scale=1)
clear_btn = gr.Button("Clear Chat", variant="secondary", scale=1)
# Example prompts
gr.Examples(
examples=[
"Hello! How are you today?",
"Tell me a joke",
"What's the weather like?",
"Can you help me with coding?",
"What's your favorite movie?",
"Explain quantum physics",
"Tell me about space exploration",
"Write a short poem about AI"
],
inputs=msg,
label="Example prompts to get started"
)
# Event handlers
msg.submit(
chat_interface,
inputs=[msg, chatbot],
outputs=[chatbot, msg]
)
send_btn.click(
chat_interface,
inputs=[msg, chatbot],
outputs=[chatbot, msg]
)
clear_btn.click(
clear_chat,
outputs=chatbot
)
return demo
if __name__ == "__main__":
# Create and launch the demo
demo = create_demo()
# Launch with optimal settings for Hugging Face Spaces
demo.launch(
share=False, # Disable sharing since this is for Hugging Face Spaces
inbrowser=False,
server_name="0.0.0.0",
server_port=7860,
show_api=False,
quiet=True
)