Added optimizer.zero_grad(), loss.backward(), and optimizer.step() properly

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AdamW
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import datetime
 import torch
 import torch.nn.functional as F
@@ -8,7 +8,9 @@ import torch.nn.functional as F
 model_name = "google/flan-t5-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-optimizer = AdamW(model.parameters(), lr=5e-5)
+
+# Define optimizer for FLAN-T5 model
+optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
 
 # Translation Models (English <-> Hindi)
 translator_en_hi = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
@@ -28,7 +30,7 @@ def translate(text, src_lang, tgt_lang):
         outputs = translator_hi_en.generate(**inputs)
         return tokenizer_hi_en.decode(outputs[0], skip_special_tokens=True)
     else:
-        return "Translation for this pair is not supported yet!"
+        return "Translation for this pair not supported yet!"
 
 # Generate Complaint Template
 def generate_complaint(issue):
@@ -38,15 +40,11 @@ def generate_complaint(issue):
 [Your Address]  
 {date}  
 To Whom It May Concern,  
-
 **Subject: Complaint Regarding {issue}**  
-
 I am writing to formally lodge a complaint regarding {issue}. The incident occurred on [Date/Location]. The specific details are as follows:  
-- **Issue:** {issue}  
-- **Evidence:** [Provide Evidence]  
-
+- Issue: {issue}  
+- Evidence: [Provide Evidence]  
 I kindly request you to take appropriate action as per the legal guidelines.  
-
 Yours sincerely,  
 [Your Name]  
 """
@@ -59,31 +57,29 @@ def compute_loss(logits, labels):
     loss = -gathered_log_probs.mean()
     return loss
 
-# Legal Query Handling with Reinforcement Learning
 def handle_legal_query(query, language):
     if language != "English":
         query = translate(query, language, "English")
 
     inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
+
+    # Generate output
     outputs = model.generate(**inputs, max_length=150)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Simple reward function - reward if response mentions key legal terms
-    reward_terms = ["law", "legal", "rights", "disability", "act", "compliance"]
-    reward = 1.0 if any(term in response.lower() for term in reward_terms) else -0.5
+    # Simple reward function (reward if response mentions legal terms)
+    reward = 1.0 if "law" in response.lower() or "legal" in response.lower() else -1.0
 
     # Compute SCST Loss
     labels = inputs['input_ids']
     logits = model(**inputs).logits
     loss = compute_loss(logits, labels)
 
-    # Adjust loss based on reward
+    # Update model weights based on reward signal
+    optimizer.zero_grad()  # Reset gradients
     loss = loss * torch.tensor(reward, dtype=torch.float)
-    
-    # Backpropagation with optimizer
-    optimizer.zero_grad()
-    loss.backward()
-    optimizer.step()
+    loss.backward()  # Backpropagation
+    optimizer.step()  # Update model weights
 
     if language != "English":
         response = translate(response, "English", language)
@@ -94,23 +90,19 @@ def handle_legal_query(query, language):
 def generate_email(issue):
     template = f"""
 Subject: Complaint Regarding {issue}
-
-Dear Sir/Madam,  
-
-I am writing to formally lodge a complaint regarding {issue}. The incident occurred on [Date/Location]. The specific details are as follows:  
-- **Issue:** {issue}  
-- **Evidence:** [Provide Evidence]  
-
-I kindly request you to take appropriate action as per the legal guidelines.  
-
+Dear Sir/Madam,
+I am writing to formally lodge a complaint regarding {issue}. The incident occurred on [Date/Location]. The specific details are as follows:
+- Issue: {issue}
+- Evidence: [Provide Evidence]
+I kindly request you to take appropriate action as per the legal guidelines.
 Yours sincerely,  
 [Your Name]
 """
     return template.strip()
 
 # Gradio Interface
-with gr.Blocks(css=".container {width: 100%; max-width: 800px;}") as app:
-    gr.Markdown("# AI Legal Assistant for Disabilities\n### Get Legal Advice and Generate Complaint Templates Instantly")
+with gr.Blocks(css=".container {width: 100%; max-width: 600px;}") as app:
+    gr.Markdown("# AI Legal Assistant for Disabilities \n### Ask legal questions and generate complaints")
 
     with gr.Row():
         query = gr.Textbox(label="Ask your legal question", placeholder="What are my rights as a disabled person?")
@@ -129,10 +121,9 @@ with gr.Blocks(css=".container {width: 100%; max-width: 800px;}") as app:
         email_btn = gr.Button("Generate Email")
         email_output = gr.Textbox(label="Generated Email", placeholder="Generated email will appear here")
 
-    # Connect functions to buttons
     submit_btn.click(handle_legal_query, inputs=[query, lang], outputs=output)
-    generate_btn.click(generate_complaint, inputs=[issue], outputs=complaint_output)
-    email_btn.click(generate_email, inputs=[issue], outputs=email_output)
+    generate_btn.click(generate_complaint, inputs=issue, outputs=complaint_output)
+    email_btn.click(generate_email, inputs=issue, outputs=email_output)
 
-# Launch app
+# Launch the app
 app.launch()