Update app.py

app.py

@@ -2,19 +2,31 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 import torch
 
-# Load a smaller model to reduce memory usage
-model_name = "distilgpt2"  # Smaller model
+# Load a model suited for code generation
+model_name = "Salesforce/codegen-350M-mono"  # This is a smaller model, choose one suited for your task
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def generate_text(input_text):
-    # Ensure input is in the correct format
-    input_tensor = tokenizer(input_text, return_tensors="pt")  # Removed clean_up_tokenization_spaces
+# Set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+def generate_code(prompt):
+    # Prepare the input for the model
+    input_tensor = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Generate code based on the prompt
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_tensor['input_ids'],
+            max_length=300,  # You can adjust this length
+            num_beams=5,     # This controls the diversity of outputs
+            early_stopping=True
+        )
     
-    # 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
+    # Decode and return the generated code
+    generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_code
 
-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 = gr.Interface(fn=generate_code, inputs="text", outputs="text", allow_flagging="never")
+iface.launch(server_name="0.0.0.0", server_port=7860)