optimise code

app.py

@@ -1,39 +1,38 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 # Load the pre-trained model and tokenizer
-
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
-model = AutoModelForCausalLM.from_pretrained(
-  "microsoft/phi-2",
-  load_in_8bit=True,
-  device_map="auto"
-)
+model_name = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 model = torch.compile(model)
 
-promptPre = f"""You are an energy-saving expert tasked to help households reduce their monthly electricity bills. 
-Given the user's appliance usage information (including device name, wattage, hours used per day, and days used per week), 
-analyze and recommend specific, practical ways they can reduce their energy consumption.
-
-Always prioritize suggestions like:
-- Reducing usage time or frequency when excessive
-- Replacing high-consumption appliances with more efficient alternatives (e.g., inverter air conditioners, LED lights)
-- Changing habits (e.g., using fans instead of air conditioners when possible)
-- Avoiding unnecessary standby power consumption
-- Considering lifestyle changes (e.g., reading books instead of watching TV)
-
-Format the output clearly with bullet points or short paragraphs.
-Be empathetic, practical, and encouraging. Focus on achievable actions for the user.
-
-Here is the user's input:
+prompt_prefix = """
+You are an energy-saving expert tasked to help households reduce their monthly electricity bills.
+Given the user's appliance usage information (device name, wattage, hours used per day, days used per week):
+1. Flag the highest energy consumers.
+2. Recommend practical, empathetic, achievable actions.
+3. Suggest appliance swaps (e.g. LED, inverter AC) and habit changes.
+Format with bullet points.
+Here is the summary:
 """
 
-def generate_recommendation(summary):
-    inputs = tokenizer(summary, return_tensors="pt").to(model.device)
+def generate_recommendation(appliance_info: str) -> str:
+    # Build the full prompt
+    prompt = prompt_prefix + appliance_info + "\n\nRecommendations:"
+    # Tokenize and move inputs to the model device
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    # Generate with no grad, limited tokens
     with torch.no_grad():
-        out = model.generate(**inputs, max_new_tokens=100, use_cache=True)
-    return tokenizer.decode(out[0], skip_special_tokens=True)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=120,
+            use_cache=True
+        )
+    # Decode and return only the recommendations section
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return text.split("Recommendations:")[-1].strip()
 
 # Define the Gradio interface
 iface = gr.Interface(
@@ -45,4 +44,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()