remove ppo training for now

app.py

@@ -1,52 +1,36 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, LogitsProcessorList, TopKLogitsWarper, TopPLogitsWarper, set_seed
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, LogitsProcessorList, TopKLogitsWarper
 import datetime
-import torch
-from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead, create_reference_model
 
 # Load FLAN-T5 for Legal Q&A
 model_name = "google/flan-t5-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(model_name)
-
-# Create a reference model for PPO
-ref_model = create_reference_model(model)
-
-# PPO Configuration
-config = PPOConfig(
-    batch_size=1,
-    learning_rate=1e-5,
-    mini_batch_size=1,
-    steps=1  # Minimal epochs
-)
-
-# Create PPO Trainer
-ppo_trainer = PPOTrainer(
-    config=config,
-    model=model,
-    ref_model=ref_model,
-    tokenizer=tokenizer
-)
-
-# 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")
+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 Function
 def translate(text, src_lang, tgt_lang):
-    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:
+    trans_model, trans_tokenizer = load_translation_model(src_lang, tgt_lang)
+    if trans_model is None:
         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):
@@ -55,47 +39,34 @@ def generate_complaint(issue):
 [Your Name]  
 [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]  
-
 I kindly request you to take appropriate action as per the legal guidelines.  
-
 Yours sincerely,  
 [Your Name]  
 """
     return template.strip()
 
-# Handle Legal Q&A with PPO
+# Handle Legal Q&A 
 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)
+    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=256)
 
-    # Logits processing using Top-K and Top-P sampling (replacement for top_k_top_p_filtering)
+    # Logits processing using Top-K sampling
     logits_processor = LogitsProcessorList([
-        TopKLogitsWarper(50),  # Top-k sampling with k=50
-        TopPLogitsWarper(0.95) # Top-p nucleus sampling with p=0.95
+        TopKLogitsWarper(50)  # Use Top-K only
     ])
 
     # Generate Response
     outputs = model.generate(**inputs, max_length=150, logits_processor=logits_processor)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Reward Signal for PPO (basic reward)
-    reward = torch.tensor([1.0]) if "legal" in response.lower() else torch.tensor([-1.0])
-
-    # PPO Step (Reinforcement Learning)
-    ppo_trainer.step([query], [outputs], [reward])
-
     if language != "English":
         response = translate(response, "English", language)
 
@@ -121,5 +92,5 @@ with gr.Blocks(css=".container {width: 100%; max-width: 600px;}") as app:
     submit_btn.click(handle_legal_query, inputs=[query, lang], outputs=output)
     generate_btn.click(generate_complaint, inputs=issue, outputs=complaint_output)
 
-# Launch the app on Hugging Face free tier
+# Launch the app
 app.launch()