Update app.py

app.py

@@ -10,9 +10,17 @@ login(token=os.getenv("HF_TOKEN"))
 
 # Load the model and tokenizer
 model_name = "ranggafermata/Fermata-v1.2-lightcoder"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+# Load tokenizer and model with eager attention
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model.eval()
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=dtype,
+    attn_implementation="eager",
+    trust_remote_code=True
+).to(device)
 
 def generate_code(prompt, max_tokens, temperature, top_p):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)