Update app.py

app.py

@@ -5,13 +5,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import json
 import os
 
-# Load the model and tokenizer from Hugging Face
-model_name = "bigcode/starcoder"  # Use StarCoder for code-related tasks
+# Load the CodeGen-2B-mono model and tokenizer from Hugging Face
+model_name = "Salesforce/codegen-2B-mono"  # Best version for CPU-friendly performance in code generation
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Ensure the model runs on CPU for Hugging Face Spaces free tier
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Ensure the model runs on CPU (important for Hugging Face Spaces free tier)
+device = torch.device("cpu")
 model.to(device)
 
 # Cache to store recent prompts and responses with file-based persistence
@@ -42,11 +42,11 @@ def code_assistant(prompt, language):
     # Tokenize the input
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     
-    # Generate response with adjusted parameters for faster responses
+    # Generate response with adjusted parameters for faster CPU response
     outputs = model.generate(
         inputs.input_ids,
         max_length=128,        # Shortened max length for quicker response
-        temperature=0.1,       # Lower temperature for more focused output
+        temperature=0.1,       # Lower temperature for focused output
         top_p=0.8,             # Slightly reduced top_p for quicker sampling
         do_sample=True
     )
@@ -73,7 +73,7 @@ iface = gr.Interface(
         gr.Dropdown(choices=["Python", "JavaScript", "Java", "C++", "HTML", "CSS", "SQL", "Other"], label="Programming Language")
     ],
     outputs="text",
-    title="CodeBand",
+    title="Code Assistant with CodeGen-2B",
     description="An AI code assistant to help you with coding queries, debugging, and code generation. Specify the programming language for more accurate responses."
 )