Update app.py

app.py

@@ -1,18 +1,24 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
 from datetime import datetime
 
-# Load Llama 3 model (quantized for CPU hosting)
-MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
+# Load Llama 3.2 (QLoRA) Model on CPU
+MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="cpu"  # Force CPU usage
+)
 
-# Load Llama Guard for content moderation
+# Load Llama Guard for content moderation on CPU
 LLAMA_GUARD_NAME = "meta-llama/Llama-Guard-3-1B-INT4"
 guard_tokenizer = AutoTokenizer.from_pretrained(LLAMA_GUARD_NAME)
-guard_model = AutoModelForCausalLM.from_pretrained(LLAMA_GUARD_NAME, torch_dtype=torch.float16, device_map="auto")
+guard_model = AutoModelForCausalLM.from_pretrained(
+    LLAMA_GUARD_NAME,
+    device_map="cpu"
+)
 
 # Define Prompt Templates
 PROMPTS = {
@@ -50,7 +56,7 @@ def moderate_input(user_input):
         return "⚠️ Content flagged by Llama Guard. Please modify your input."
     return None  # Safe input, proceed normally
 
-# Function: Generate AI responses (Project Analysis, Code, or Risks)
+# Function: Generate AI responses
 def generate_response(prompt_type, **kwargs):
     prompt = PROMPTS[prompt_type].format(**kwargs)
     
@@ -58,11 +64,11 @@ def generate_response(prompt_type, **kwargs):
     if moderation_warning:
         return moderation_warning  # Stop processing if flagged
 
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=2048, truncation=True)
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
     
     outputs = model.generate(
         inputs.input_ids,
-        max_length=2048,
+        max_length=1024,
         temperature=0.7 if prompt_type == "project_analysis" else 0.5,
         top_p=0.9
     )
@@ -133,7 +139,7 @@ def create_gradio_interface():
                 AI:"""
                 
                 inputs = tokenizer(prompt, return_tensors="pt")
-                outputs = model.generate(inputs.input_ids, max_length=2048)
+                outputs = model.generate(inputs.input_ids, max_length=1024)
                 response = tokenizer.decode(outputs[0], skip_special_tokens=True)
                 chat_history.append((message, response))
                 return "", chat_history