llama / app.py
hereoncollab's picture
Update app.py
751924a verified
raw
history blame
647 Bytes
import gradio as gr
from transformers import pipeline
# Use a smaller model for lower memory usage
pipe = pipeline("text-generation", model="distilgpt2") # DistilGPT-2 is smaller and more efficient
def generate_response(prompt):
responses = pipe(prompt, max_length=50, num_return_sequences=1)
return responses[0]['generated_text']
interface = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(label="Enter your message:", lines=2, placeholder="Type your message here..."),
outputs="text",
title="Smaller Model Chatbot",
description="A chatbot using a smaller model for lower memory usage."
)
interface.launch()