Update app.py

app.py

@@ -1,28 +1,28 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Load DeepSeek model
-deepseek_tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1")
-deepseek_model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1", torch_dtype="auto", device_map="auto")
+# Load DeepSeek with trust_remote_code
+deepseek_tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
+deepseek_model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True, torch_dtype="auto", device_map="auto")
 deepseek_pipe = pipeline("text-generation", model=deepseek_model, tokenizer=deepseek_tokenizer)
 
-# Load LLaMA model
-llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-4-Scout-17B-16E-Instruct")
-llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-4-Scout-17B-16E-Instruct", torch_dtype="auto", device_map="auto")
+# Load LLaMA with trust_remote_code
+llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-4-Scout-17B-16E-Instruct", trust_remote_code=True)
+llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-4-Scout-17B-16E-Instruct", trust_remote_code=True, torch_dtype="auto", device_map="auto")
 llama_pipe = pipeline("text-generation", model=llama_model, tokenizer=llama_tokenizer)
 
 def generate_and_enhance_code(code_request: str, features: str):
-    # Step 1: Generate base code from DeepSeek
+    # Generate base code from DeepSeek
     base_output = deepseek_pipe(code_request, max_new_tokens=512, do_sample=True, temperature=0.7)[0]["generated_text"]
 
-    # Step 2: Ask Llama to add features
+    # Enhance with LLaMA
     enhancement_prompt = f"Hey Llama! can you please add some more features in my code?\n\nOriginal code:\n{base_output}\n\nFeatures to add:\n{features}\n\nAdd the features and pass me the code without any extra asking!"
     enhanced_output = llama_pipe(enhancement_prompt, max_new_tokens=1024, do_sample=True, temperature=0.6)[0]["generated_text"]
 
     return enhanced_output
 
 with gr.Blocks() as demo:
-    gr.Markdown("## AI Code Assistant with DeepSeek + LLaMA 4")
+    gr.Markdown("## MINEOGO: DeepSeek + LLaMA Code Assistant")
     with gr.Row():
         code_input = gr.Textbox(lines=5, label="What code do you want?")
         feature_input = gr.Textbox(lines=3, label="What features should LLaMA add?")