Update app.py

app.py

@@ -1,65 +1,56 @@
-# Copyright (c) 2008-2021, Hazelcast, Inc. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Model configuration
-model_name = "Salesforce/codet5p-770m"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# Model configuration: CodeGemma 1.1B
+MODEL_NAME = "google/codegemma-1.1b"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Prompts for different languages
-language_prompts = {
-    "Python": "Fix this Python code:\n",
-    "C": "Fix this C code:\n",
-    "C++": "Fix this C++ code:\n",
-    "JavaScript": "Fix this JavaScript code:\n"
+# Supported languages and their prompt templates
+LANGUAGE_PROMPTS = {
+    "Python": "Fix this Python code:",
+    "C": "Fix this C code:",
+    "C++": "Fix this C++ code:",
+    "JavaScript": "Fix this JavaScript code:"
 }
 
-# Debugging logic
 def eternos_debugger(code, error, language):
     if not code.strip():
         return "❌ Please provide code."
-
-    prompt = f"{language_prompts[language]}{code}\nError:\n{error}\nCorrected code:\n"
+    prompt = f"{LANGUAGE_PROMPTS[language]}\n{code}\nError:\n{error}\nCorrected code:\n"
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
     outputs = model.generate(
         **inputs,
         max_new_tokens=256,
-        temperature=0.2,
-        do_sample=False
+        temperature=0.1,
+        top_p=0.9,
+        do_sample=False,
+        pad_token_id=tokenizer.eos_token_id
     )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response.strip()
 
-# UI setup
-with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {background-color: #cbedec;}") as demo:
-    gr.Markdown("## ⚙️ Eternos — AI Code Debugger")
-    gr.Markdown("Supports Python, C, C++, JavaScript — powered by CodeT5p")
+# Gradio user interface
+def create_interface():
+    with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {background-color: #cbedec;}") as demo:
+        gr.Markdown("## ⚙️ Eternos — AI Code Debugger")
+        gr.Markdown("Supports Python, C, C++, JavaScript — powered by CodeGemma 1.1B")
+
+        with gr.Row():
+            code_input = gr.Textbox(label="📝 Your Code", lines=12)
+            error_input = gr.Textbox(label="⚠️ Error Message (optional)", lines=4)
 
-    with gr.Row():
-        code_input = gr.Textbox(label="📝 Your Code", lines=12)
-        error_input = gr.Textbox(label="⚠️ Error Message (optional)", lines=4)
+        language_input = gr.Dropdown(["Python", "C", "C++", "JavaScript"], label="🌐 Language", value="Python")
+        output_code = gr.Code(label="✅ Suggested Fix")
+        run_btn = gr.Button("🛠️ Fix Code")
 
-    language_input = gr.Dropdown(["Python", "C", "C++", "JavaScript"], label="🌐 Language", value="Python")
-    output_code = gr.Code(label="✅ Suggested Fix")
-    run_btn = gr.Button("🛠️ Fix Code")
+        run_btn.click(fn=eternos_debugger, inputs=[code_input, error_input, language_input], outputs=output_code)
 
-    run_btn.click(fn=eternos_debugger, inputs=[code_input, error_input, language_input], outputs=output_code)
+    return demo
 
-demo.launch()
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()