Update app.py

app.py

@@ -1,19 +1,23 @@
 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']
+pipe = pipeline("text-generation", model="distilgpt2")
+
+def generate_response(user_input):
+    formatted_input = f"Human: {user_input}\nAI:"
+    response = pipe(formatted_input, max_length=150, num_return_sequences=1)
+    generated_text = response[0]['generated_text']
+    ai_response = generated_text.split('AI:')[-1].strip()
+    return ai_response
+
 
 interface = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(label="Enter your message:", lines=2, placeholder="Type your message here..."),
+    inputs=gr.Textbox(label="shitty mesage:", lines=2, placeholder="Type your message here..."),
     outputs="text",
-    title="Smaller Model Chatbot",
-    description="A chatbot using a smaller model for lower memory usage."
+    title="distilgpt2",
+    description="distel gppt 20"
 )
 
 interface.launch()