Skip to content

feat: Implement self correcting CHIME model and federated learning#79

Merged
pradeeban merged 30 commits intoKathiraveluLab:devfrom
anish1206:exp/bert-improv
Feb 1, 2026
Merged

feat: Implement self correcting CHIME model and federated learning#79
pradeeban merged 30 commits intoKathiraveluLab:devfrom
anish1206:exp/bert-improv

Conversation

@anish1206
Copy link
Contributor

This PR introduces a Self-Correcting AI System for the DREAMS classification model (CHIME). It implements a closed-loop Federated Learning workflow where user corrections to sentiment predictions are used to improve the model's accuracy automatically and privately, without centralized data aggregation.

Key Features Implemented

  • User Feedback Loop:

    • Added "Edit Prediction" workflow in the dashboard.
    • Users can correct the AI's classification (Connectedness, Hope, Identity, Meaning, Empowerment) or flag posts as "None".
  • Database Schema Updates:

    • Updated Post model to store corrected_label and is_fl_processed status.
  • Federated Learning Worker (fl_worker.py):

    • Continuous Learning: Loads the previous best model state instead of starting from scratch.
    • Siloed Training: Retrains only the classifier head (freezes BERT base) on user corrections.
    • Validation Gate: Automatically rejects updates that fail "Anchor Tests" (preventing catastrophic forgetting) or fail to learn the new data.
    • Atomic Swaps: Updates the live production model only if validation passes.
  • Event Driven Automation:

    • Automatically triggers the training worker in a background thread once BATCH_SIZE (e.g. 50) corrections accumulate.
    • Zero latency for the user (non-blocking execution).
  • Smart Inference:

    • Updated sentiment.py to prioritize the locally learned model (production_chime_model) over the generic Hugging Face base model.

Architecture

image

How Corrections are Handled

  1. Threshold: Training triggers at 50 corrections.
  2. Validation: The new model must maintain high accuracy on a static "Anchor Set" of clear examples.
  3. Fallback: If the local model is missing or corrupted, the system gracefully falls back to the original Huggingface model.

Impact

This feature transforms the model from a static classifier into an adaptive, user aligned system. By implementing this architecture, we ensure that the model evolve with our users' actual lived experiences, rather than relying solely on pre trained external datasets. This is achieved with zero compromise on user privacy or system performance.

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @anish1206, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant enhancement to the DREAMS application by integrating a self-correcting AI system for the CHIME classification model. It establishes a robust Federated Learning pipeline that leverages user feedback to continuously refine the model's predictions. The system is designed for privacy, efficiency, and stability, ensuring that the model adapts to real-world data while maintaining performance and preventing degradation.

Highlights

  • Self-Correcting AI System: Implemented a closed-loop Federated Learning workflow for the DREAMS CHIME classification model, allowing it to improve accuracy based on user corrections without centralized data aggregation.
  • User Feedback Loop: Added an 'Edit Prediction' workflow in the dashboard, enabling users to correct AI classifications (Connectedness, Hope, Identity, Meaning, Empowerment) or flag posts as 'None'.
  • Federated Learning Worker (fl_worker.py): Introduced a dedicated worker for continuous learning, which loads the previous best model state, retrains only the classifier head (freezing BERT base), and includes a validation gate to prevent catastrophic forgetting or failed learning.
  • Event-Driven Automation: Training is automatically triggered in a background thread once a BATCH_SIZE (e.g., 50) of corrections accumulate, ensuring zero latency for the user.
  • Smart Inference: The sentiment.py module now prioritizes loading the locally learned model (production_chime_model) over the generic Hugging Face base model for predictions.
  • Database Schema Updates: The Post model was updated to store corrected_label and is_fl_processed status, facilitating the tracking of user feedback and training progress.
  • Comprehensive Documentation and Testing: Added detailed documentation (federated-learning.md) explaining the architecture and functionality, along with an end-to-end test script (test_fl.py) for the FL loop.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a robust self-correcting AI system for the CHIME model using federated learning, which is a significant enhancement. The implementation includes a user feedback loop, database schema updates, a dedicated federated learning worker, and smart inference. The documentation is comprehensive and clearly explains the architecture and mechanisms. Several areas for improvement have been identified, primarily concerning error handling, configuration management, and consistency in logging and testing. Addressing these will further enhance the system's reliability and maintainability.

anish1206 and others added 10 commits January 30, 2026 12:06
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@pradeeban
Copy link
Member

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces an impressive self-correcting mechanism for the CHIME model using a federated learning approach. The implementation is well-structured, with a clear separation of concerns between the web endpoint, the background worker, and the model loading logic. The inclusion of a validation gate with anchor tests to prevent catastrophic forgetting is a great safety feature. The addition of comprehensive documentation and an end-to-end test script is also highly commendable.

My review focuses on a critical security vulnerability in the feedback submission endpoint, the robustness of the model update process, and the reliability of the new test case. Addressing these points will significantly improve the security and stability of this new feature.

Comment on lines 188 to 194
if os.path.exists(PRODUCTION_MODEL_DIR):
shutil.rmtree(PRODUCTION_MODEL_DIR)

# Ensure parent dict exists
os.makedirs(os.path.dirname(PRODUCTION_MODEL_DIR), exist_ok=True)

shutil.copytree(TEMP_MODEL_DIR, PRODUCTION_MODEL_DIR)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current model update process (rmtree then copytree) is not atomic. If the script fails after deleting the old model but before the new one is fully copied, the application will be left in a broken state with no model. For a truly atomic update, you should use directory renames. A common pattern is:

  1. Save the new model to a temporary directory.
  2. Rename the current production directory to a backup name.
  3. Rename the temporary directory to the production name.
  4. Clean up the backup directory.

from dreamsApp.app import create_app
from dreamsApp.app.fl_worker import run_federated_round

def test_fl_loop():
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This test is brittle because it depends on the BATCH_SIZE constant from fl_worker.py. The test creates 10 mock posts, but the production BATCH_SIZE is 50. As a result, the training logic in run_federated_round won't be triggered, and the test will fail.

To make the test robust, you could either:

  1. Update the test to create 50 mock posts to meet the threshold.
  2. Use a mocking library (like unittest.mock.patch) to temporarily set fl_worker.BATCH_SIZE to 10 for this test case.


# Update the post using $set to add correction data
result = mongo.update_one(
{'_id': ObjectId(post_id), 'user_id': data.get('user_id')},
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The call ObjectId(post_id) will raise a bson.errors.InvalidId exception if the post_id from the request is malformed, which will cause an unhandled server error (500). It's better to catch this specific exception and return a 400 Bad Request, as it's a client-side error.

pradeeban and others added 2 commits January 31, 2026 10:41
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@pradeeban
Copy link
Member

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a robust self-correcting AI system for the DREAMS classification model using Federated Learning. Key features include a user feedback loop, database schema updates, a dedicated Federated Learning worker with continuous learning and validation gates, event-driven automation, and smart inference prioritizing the locally learned model. The changes are well-documented and include a comprehensive new documentation file (federated-learning.md).

Overall, the implementation is well-structured and addresses the core requirements effectively. I've identified a few areas for improvement related to code cleanliness, logging consistency, and ensuring test accuracy.


<script>
const userId = "{{ user_id }}"; // make sure this is being passed in the Flask route
const userId = "{{ user_id }}";
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The userId constant is declared twice in the JavaScript. It's already declared earlier in the script. It's best to declare constants only once to avoid redundancy and potential issues.

pradeeban and others added 3 commits January 31, 2026 10:46
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@pradeeban
Copy link
Member

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a significant and well-designed feature for a self-correcting AI model using a federated learning approach. The implementation is comprehensive, covering the user feedback loop, a background training worker, safety validation gates, and excellent documentation. My review highlights a critical race condition in the training trigger mechanism that could lead to data inconsistencies and wasted resources. I've also included suggestions to improve code style, security, and the correctness of the new tests.

anish1206 and others added 2 commits February 1, 2026 12:31
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
…ing.

Add audit trail, stale-lock/processing cleanup, improved logging and tests.
@anish1206
Copy link
Contributor Author

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a well-designed, self-correcting AI system using federated learning. The implementation is robust, featuring crucial safety mechanisms like validation gates, atomic model swaps, and database locks to prevent race conditions. The code is generally clean and includes excellent documentation and tests. My feedback focuses on minor improvements for consistency, maintainability, and documentation accuracy.

anish1206 and others added 2 commits February 1, 2026 14:15
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@anish1206
Copy link
Contributor Author

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This is an excellent pull request that introduces a robust, self-correcting mechanism for the CHIME model using federated learning principles. The implementation is well-structured, featuring a non-blocking, event-driven architecture, safety gates to prevent model degradation, and comprehensive logging and documentation. The use of a distributed lock for training runs is a great touch for ensuring stability.

My review focuses on further improving the robustness and maintainability of the system. Key suggestions include:

  • Implementing a truly atomic model swap to prevent data loss during deployment.
  • Making timestamp generation timezone-aware.
  • Improving the flexibility of the model layer freezing logic.
  • Ensuring configuration consistency between different parts of the application.

Overall, this is a very strong contribution that significantly enhances the application's capabilities.

anish1206 and others added 5 commits February 1, 2026 15:55
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@anish1206
Copy link
Contributor Author

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This is an impressive and well-architected pull request that introduces a self-correcting AI system for the CHIME model using a federated learning approach. The implementation is robust, featuring a user feedback loop, a background training worker with distributed locking, validation gates to prevent model degradation, and atomic model swaps for safe deployment. The code is clear, well-documented, and includes comprehensive tests. My review includes a few suggestions to further improve consistency in timestamp handling, enhance the robustness of the training logic, and refine the user experience in the feedback submission process.


# 3. Training Loop
model.train()
optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=LEARNING_RATE)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The optimizer is hardcoded to use model.classifier.parameters(). While this works for many standard Hugging Face sequence classification models, it's not robust. If a model with a differently named classification head is used in the future, this will break. A more resilient approach is to pass all parameters that require gradients to the optimizer.

Suggested change
optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=LEARNING_RATE)
optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=LEARNING_RATE)

Comment on lines +272 to +273
alert('Input recorded! The model will learn from this.');
location.reload();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using alert() and location.reload() creates a disruptive user experience. The alert dialog blocks the UI, and the full page reload is jarring. A better approach would be to update the UI dynamically upon a successful API call. You could replace the buttons with a "Thank you" message and update the classification label and badge directly using JavaScript, providing instant and seamless feedback to the user.

anish1206 and others added 2 commits February 1, 2026 16:10
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@anish1206
Copy link
Contributor Author

Hey @pradeeban , I believe most of the critical/high priority issues have been resolved.

could you please take a look over the changes and tell me if any more updates are required.

@pradeeban
Copy link
Member

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This is an excellent pull request that introduces a well-architected self-correcting AI system using federated learning. The implementation demonstrates strong engineering practices, including the use of atomic operations for database locking and model deployment, robust validation gates to prevent model degradation, and comprehensive documentation and testing. My feedback is focused on minor but important improvements for production-readiness: enhancing efficiency in data handling and standardizing the use of UTC for all timestamp operations to prevent potential timezone-related bugs.

pradeeban and others added 3 commits February 1, 2026 09:18
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@pradeeban
Copy link
Member

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a robust and well-designed self-correcting system for the CHIME model using a federated learning approach. The implementation is excellent, covering the entire workflow from user feedback collection on the frontend, a secure and rate-limited API endpoint, to a sophisticated background worker for model training. The worker includes crucial safety features like a validation gate with anchor tests to prevent catastrophic forgetting, and an atomic swap mechanism for safe model deployment. The addition of comprehensive documentation and an end-to-end test suite is commendable and greatly enhances the maintainability of this new feature. My review includes a few suggestions for improvement in the fl_worker.py file, mainly around making hyperparameters and validation thresholds configurable to increase flexibility, and a note on training scalability for future consideration. Overall, this is a high-quality contribution.

# dreamsApp/app/models/temp_training_artifact
TEMP_MODEL_DIR = os.path.join(BASE_DIR, "models", "temp_training_artifact")

BATCH_SIZE = 50
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The global BATCH_SIZE constant is redundant as BATCH_SIZE is fetched from the app config on line 113 within run_federated_round. To avoid confusion, it's best to remove this global constant and rely solely on the application configuration.

TEMP_MODEL_DIR = os.path.join(BASE_DIR, "models", "temp_training_artifact")

BATCH_SIZE = 50
LEARNING_RATE = 1e-5 # Conservative learning rate
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The LEARNING_RATE (here) and EPOCHS (on line 204) are hardcoded. It's a good practice to make these hyperparameters configurable through the application config (e.g., app.config.get(...)). This allows for easier tuning of the training process without modifying the source code.

logger.debug(f"[Anchor Fail] Text: '{example['text'][:30]}...' Expected: {target_str}, Got: {pred_str}")

logger.info(f"[Safety Check] Anchor Accuracy: {correct_anchors}/{len(ANCHOR_EXAMPLES)}")
if correct_anchors < 4: # Stricter check for catastrophic forgetting
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The validation threshold for the anchor check is hardcoded (< 4). The improvement check threshold on line 81 is also hardcoded (< 0.5). Making these values configurable (e.g., MIN_ANCHOR_ACCURACY, MIN_IMPROVEMENT_ACCURACY) via the application config would provide more flexibility to adjust the strictness of the validation gate as the model evolves.

Comment on lines +200 to +211
texts = [item[0] for item in training_data]
labels_tensor = torch.tensor([item[1] for item in training_data])
inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")

EPOCHS = 3
for epoch in range(EPOCHS):
optimizer.zero_grad()
outputs = model(**inputs, labels=labels_tensor)
loss = outputs.loss
loss.backward()
optimizer.step()
logger.info(f"[Epoch {epoch+1}/{EPOCHS}] Loss: {loss.item():.4f}")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The training loop processes the entire batch of training_data at once. While this works for the current batch size of 50, it could lead to out-of-memory issues if the batch size is increased significantly. For future scalability, consider refactoring this to use a torch.utils.data.DataLoader to process the data in smaller mini-batches within each epoch.

@pradeeban pradeeban merged commit d9cc021 into KathiraveluLab:dev Feb 1, 2026
1 check passed
KrishanYadav333 added a commit to KrishanYadav333/DREAMS that referenced this pull request Feb 6, 2026
…tation, code stubs, and tests

This PR establishes the foundation for multi-dimensional location-proximity analysis
in DREAMS, building upon existing EXIF extraction (PR KathiraveluLab#77) and emotion proximity (PR KathiraveluLab#70).

## Documentation (9 new/updated files)
- docs/api_design.md: REST API specification for location-proximity endpoints
- docs/evaluation_metrics.md: Quantitative metrics and ablation study plan
- docs/exif_extraction_research.md: Library comparison research (informed PR KathiraveluLab#77)
- docs/integration_guide.md: Step-by-step integration instructions
- docs/project_roadmap.md: GSoC 2026 timeline aligned with official dates (350h)
- docs/risk_analysis.md: Risk matrix and mitigation strategies
- docs/TEST_PLAN.md: Extended with 50+ location-proximity test cases
- plans/pre_gsoc_contribution_plan.md: 7-week, 18-PR contribution roadmap
- dreamsApp/docs/data-model.md: Added location_analysis and emotion_location_entries collections

## Code Implementation
- dreamsApp/exif_extractor.py: NEW - Complete EXIF extraction with dual-library fallback
- dreamsApp/location_proximity.py: Updated stubs with EXIFExtractor integration
- ARCHITECTURE.md: Updated diagram to show integration with PR KathiraveluLab#77 and KathiraveluLab#70
- LOCATION_PROXIMITY_SUMMARY.md: Added acknowledgment of existing work

## Tests
- tests/test_exif_extraction.py: NEW - Unit tests for EXIF extractor with mocking

## Code Quality
- Removed emojis from entire project (8 files) for professional documentation
- data_integrity/reporter.py: Replaced emoji indicators with text
- dream-integration/app/templates/index.html: Replaced emoji UI elements
- dreamsApp/app/dashboard/main.py: Removed emoji comments

## Integration Points
- Builds upon PR KathiraveluLab#77 (kunal-595): Uses existing EXIFExtractor class
- Complements PR KathiraveluLab#70 (AnvayKharb): Adds spatial proximity to time-aware emotion analysis
- Aligns with PR KathiraveluLab#79 (anish1206): Emotion-location work supports CHIME framework

## Key Features
- Multi-dimensional proximity: geographic + categorical + linguistic + cultural
- Emotion-location hotspot detection
- Semantic clustering with DBSCAN
- MongoDB schema extensions for location data
- Performance benchmarks and evaluation framework

Total: 10 new files, 9 updated files, 350 hours planned for GSoC 2026 implementation
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants