Update app.py

app.py

@@ -2,17 +2,19 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 import torch
 
-# Load the model
-model_name = "HuggingFaceH4/starchat2-15b-v0.1"  # Your main model
+# Load a smaller model to reduce memory usage
+model_name = "distilgpt2"  # Smaller model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
 
 def generate_text(input_text):
     # Ensure input is in the correct format
     input_tensor = tokenizer(input_text, return_tensors="pt", clean_up_tokenization_spaces=True)
-    output = model.generate(**input_tensor)
+    
+    # Generate text with a limit on max_length to reduce memory usage
+    output = model.generate(**input_tensor, max_length=50)  # Adjust max_length as needed
     response = tokenizer.decode(output[0], skip_special_tokens=True)
     return response
 
 iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", allow_flagging="never")
-iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch(server_name="0.0.0.0", server_port=7860)