Update app.py

app.py

@@ -1,3 +1,37 @@
+import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 
-gr.load("models/rootxhacker/CodeAstra-7B").launch()
+# Load the model and tokenizer
+peft_model_id = "rootxhacker/CodeAstra-7B"
+config = PeftConfig.from_pretrained(peft_model_id)
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_4bit=True, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+
+# Load the Lora model
+model = PeftModel.from_pretrained(model, peft_model_id)
+
+model.config.to_json_file("config.json")
+
+def get_completion(query, model, tokenizer):
+    inputs = tokenizer(query, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def code_review(code_to_analyze):
+    query = f"As a code review expert, your role will be to carefully examine the code for potential security flaws and provide guidance on secure coding practices. This may include identifying common coding mistakes that could lead to vulnerabilities, suggesting ways to improve the code's overall security, and recommending tools or techniques that can be used to detect and prevent potential threats. Your expertise in security will be particularly valuable in ensuring that any code developed meets the highest security standard:\n{code_to_analyze}"
+    result = get_completion(query, model, tokenizer)
+    return result
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=code_review,
+    inputs=gr.Textbox(lines=10, label="Enter code to analyze"),
+    outputs=gr.Textbox(label="Code Review Result"),
+    title="Code Review Expert",
+    description="This tool analyzes code for potential security flaws and provides guidance on secure coding practices."
+)
+
+# Launch the Gradio app with a public link
+iface.launch(share=True)