Update app.py

app.py

@@ -1,19 +1,30 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Load your model from Hugging Face Hub
-model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free"
+# Load the pre-trained DistilGPT-2 model and tokenizer
+model_name = "distilgpt2"  # This is a smaller version of GPT-2
 model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Function to generate text from the model
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_length=50)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# Function to generate text based on user input
+def generate_response(user_input):
+    # Encode the input prompt
+    inputs = tokenizer.encode(user_input, return_tensors="pt")
 
-# Set up Gradio Interface
-iface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
+    # Generate a response from the model
+    with torch.no_grad():  # Disable gradient calculation for inference
+        outputs = model.generate(inputs, max_length=100, num_return_sequences=1, no_repeat_ngram_size=2, top_p=0.9, top_k=50)
+    
+    # Decode the generated response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return response
 
-# Launch the app
+# Create Gradio interface for text input and output
+iface = gr.Interface(fn=generate_response, inputs="text", outputs="text", 
+                     title="DistilGPT-2 Chatbot", 
+                     description="A lightweight conversational chatbot using DistilGPT-2.")
+
+# Launch the Gradio app
 iface.launch()