Use SCST RLAI and check

app.py

@@ -1,36 +1,33 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, LogitsProcessorList, TopKLogitsWarper
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import datetime
+import torch
+import torch.nn.functional as F
 
 # Load FLAN-T5 for Legal Q&A
 model_name = "google/flan-t5-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Translation Models (Dynamically Loaded)
-translation_models = {
-    "en-hi": ("Helsinki-NLP/opus-mt-en-hi", "Helsinki-NLP/opus-mt-en-hi"),
-    "hi-en": ("Helsinki-NLP/opus-mt-hi-en", "Helsinki-NLP/opus-mt-hi-en")
-}
-
-def load_translation_model(src_lang, tgt_lang):
-    pair = f"{src_lang[:2]}-{tgt_lang[:2]}"
-    if pair in translation_models:
-        model_name, tokenizer_name = translation_models[pair]
-        trans_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        trans_tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-        return trans_model, trans_tokenizer
-    return None, None
+# Translation Models (English <-> Hindi)
+translator_en_hi = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
+tokenizer_en_hi = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
+
+translator_hi_en = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-hi-en")
+tokenizer_hi_en = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-hi-en")
 
 # Translation Function
 def translate(text, src_lang, tgt_lang):
-    trans_model, trans_tokenizer = load_translation_model(src_lang, tgt_lang)
-    if trans_model is None:
+    if src_lang == "English" and tgt_lang == "Hindi":
+        inputs = tokenizer_en_hi(text, return_tensors="pt", padding=True, truncation=True)
+        outputs = translator_en_hi.generate(**inputs)
+        return tokenizer_en_hi.decode(outputs[0], skip_special_tokens=True)
+    elif src_lang == "Hindi" and tgt_lang == "English":
+        inputs = tokenizer_hi_en(text, return_tensors="pt", padding=True, truncation=True)
+        outputs = translator_hi_en.generate(**inputs)
+        return tokenizer_hi_en.decode(outputs[0], skip_special_tokens=True)
+    else:
         return "Translation for this pair not supported yet!"
-    
-    inputs = trans_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    outputs = trans_model.generate(**inputs, max_length=256)
-    return trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Generate Complaint Template
 def generate_complaint(issue):
@@ -50,31 +47,65 @@ Yours sincerely,
 """
     return template.strip()
 
-# Handle Legal Q&A 
+# Self-Critical Sequence Training (SCST) for RL
+def compute_loss(logits, labels):
+    log_probs = F.log_softmax(logits, dim=-1)
+    gathered_log_probs = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
+    loss = -gathered_log_probs.mean()
+    return loss
+
 def handle_legal_query(query, language):
     if language != "English":
         query = translate(query, language, "English")
 
-    # Tokenize input
-    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=256)
-
-    # Logits processing using Top-K sampling
-    logits_processor = LogitsProcessorList([
-        TopKLogitsWarper(50)  # Use Top-K only
-    ])
+    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
 
-    # Generate Response
-    outputs = model.generate(**inputs, max_length=150, logits_processor=logits_processor)
+    # 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 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)
+
+    # Update model weights based on reward signal
+    loss = loss * torch.tensor(reward, dtype=torch.float)
+    loss.backward()
+    model.optimizer.step()
+    model.zero_grad()
+
     if language != "English":
         response = translate(response, "English", language)
 
     return response
 
-# Define Gradio Interface
+# Generate Email
+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.
+
+Yours sincerely,  
+[Your Name]
+"""
+    return template.strip()
+
+# Gradio Interface
 with gr.Blocks(css=".container {width: 100%; max-width: 600px;}") as app:
-    gr.Markdown("# AI Legal Assistant\n### Ask legal questions and generate complaints")
+    gr.Markdown("# AI Legal Assistant for disabilities
+### 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?")
@@ -89,8 +120,12 @@ with gr.Blocks(css=".container {width: 100%; max-width: 600px;}") as app:
         generate_btn = gr.Button("Generate Complaint")
         complaint_output = gr.Textbox(label="Generated Complaint", placeholder="Complaint template will appear here")
 
+    with gr.Row():
+        email_btn = gr.Button("Generate Email")
+        email_output = gr.Textbox(label="Generated Email", placeholder="Generated email will appear here")
+
     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)
 
-# Launch the app
 app.launch()