Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| from transformers import AutoModel, AutoTokenizer | |
| from peft import LoraConfig, get_peft_model | |
| import numpy as np | |
| import os | |
| from huggingface_hub import hf_hub_download | |
| device = torch.device('cpu') # HF Free Tier uses CPU | |
| # DOWNLOAD THE HUGE MODEL FROM YOUR MODEL REPO INSTEAD OF LOCAL STORAGE | |
| print("Downloading heavy model weights from hnninioi/AyurGenixV9-Model...") | |
| model_path = hf_hub_download(repo_id="hnninioi/AyurGenixV9-Model", filename="AyurGenixV9_FULLY_EMBEDDED.pt") | |
| ckpt = torch.load(model_path, map_location=device, weights_only=False) | |
| cfg = ckpt['config'] | |
| all_herbs = ckpt['all_herbs'] | |
| raw_herb_props = ckpt.get('herb_properties', {}) | |
| if isinstance(raw_herb_props, list): | |
| herb_props = {str(h.get('name', '')).lower(): h for h in raw_herb_props} | |
| else: | |
| herb_props = raw_herb_props | |
| formulations = ckpt.get('formulations', []) | |
| interaction_model = ckpt.get('interaction_model', None) | |
| le_dosha = ckpt['label_encoder_dosha'] | |
| side_effect_keywords = ckpt.get('side_effect_keywords', []) | |
| CLASSICAL_INTERACTIONS = ckpt.get('classical_interactions', {}) | |
| HERB_ALIASES = ckpt.get('herb_aliases', {}) | |
| HERB_ALIASES['licorice'] = 'yashtimadhu' | |
| HERB_ALIASES['mulethi'] = 'yashtimadhu' | |
| class AyurGenixV8(nn.Module): | |
| def __init__(self, encoder, hidden_size, num_herbs, num_doshas, num_severity, num_se): | |
| super().__init__() | |
| self.encoder = encoder | |
| h = hidden_size | |
| self.trunk = nn.Sequential(nn.Linear(h, 512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(0.2)) | |
| self.herb_head = nn.Sequential(nn.Linear(512, 256), nn.LayerNorm(256), nn.GELU(), nn.Dropout(0.1), nn.Linear(256, num_herbs)) | |
| self.dosha_head = nn.Sequential(nn.Linear(512, 128), nn.GELU(), nn.Linear(128, num_doshas)) | |
| self.severity_head = nn.Sequential(nn.Linear(512, 128), nn.GELU(), nn.Linear(128, num_severity)) | |
| self.conflict_head = nn.Sequential(nn.Linear(512, 64), nn.GELU(), nn.Linear(64, 1)) | |
| self.toxicity_head = nn.Sequential(nn.Linear(512, 64), nn.GELU(), nn.Linear(64, 3)) | |
| self.side_effect_head = nn.Sequential(nn.Linear(512, 128), nn.GELU(), nn.Linear(128, num_se)) | |
| self.dosage_head = nn.Sequential(nn.Linear(512, 64), nn.GELU(), nn.Linear(64, 1)) | |
| def forward(self, input_ids, attention_mask): | |
| cls = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :] | |
| shared = self.trunk(cls) | |
| return {'herb': self.herb_head(shared), 'dosha': self.dosha_head(shared), 'severity': self.severity_head(shared), 'conflict': self.conflict_head(shared).squeeze(-1), 'toxicity': self.toxicity_head(shared), 'side_effects': self.side_effect_head(shared), 'dosage': self.dosage_head(shared).squeeze(-1)} | |
| model_name = cfg.get('backbone', 'ai4bharat/IndicBERTv2-MLM-only') | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| base_encoder = AutoModel.from_pretrained(model_name) | |
| lora_encoder = get_peft_model(base_encoder, LoraConfig(r=cfg.get('lora_r', 16), lora_alpha=cfg.get('lora_alpha', 32), lora_dropout=0.1, target_modules=cfg.get('lora_targets', ['query', 'value', 'key']), bias='none')) | |
| model = AyurGenixV8(lora_encoder, cfg['hidden_size'], cfg['num_herbs'], cfg['num_doshas'], cfg['num_severity'], cfg['num_side_effects']).to(device) | |
| model.load_state_dict(ckpt['v8_model_state_dict']) | |
| model.eval() | |
| SEV_NAMES = ['Mild', 'Mild-Moderate', 'Moderate', 'Moderate-Severe', 'Severe'] | |
| TOX_NAMES = ['Low', 'Medium', 'High'] | |
| LABEL_NAMES = ['SYNERGISTIC', 'CAUTION', 'CONTRAINDICATED'] | |
| VIRYA_CONFLICT = {'ushna': 0, 'sheeta': 1, 'hot': 0, 'cold': 1} | |
| RASA_GROUPS = {'madhura': 0, 'sweet': 0, 'amla': 1, 'lavana': 2, 'tikta': 3, 'katu': 4, 'kashaya': 5} | |
| def resolve_herb_name(name): | |
| name = str(name).lower().strip() | |
| if name in herb_props: return name | |
| if name in HERB_ALIASES: | |
| if HERB_ALIASES[name] in herb_props: return HERB_ALIASES[name] | |
| for db_name in herb_props: | |
| if name in db_name or db_name in name: return db_name | |
| return None | |
| def tokenize_indications(text): | |
| keys = ['kasa', 'cough', 'jwara', 'fever', 'amavata', 'joint', 'arthritis', 'indigestion', 'prameha', 'diabetes', 'liver', 'skin', 'shwasa', 'asthma'] | |
| return {k for k in keys if k in text.lower()} | |
| def encode_herb(props): | |
| features = [] | |
| virya_val = -1 | |
| for key, val in VIRYA_CONFLICT.items(): | |
| if key in props.get('virya', ''): | |
| virya_val = val | |
| break | |
| features.append(virya_val) | |
| rasa_vec = [0] * 6 | |
| for key, idx in RASA_GROUPS.items(): | |
| if key in props.get('rasa', ''): rasa_vec[idx] = 1 | |
| features.extend(rasa_vec) | |
| dosha_text = props.get('dosha_effect', '') | |
| features.extend([int('vata' in dosha_text), int('pitta' in dosha_text), int('kapha' in dosha_text)]) | |
| return np.array(features, dtype=float) | |
| def rule_based_reason(props_a, props_b, label_idx): | |
| va, vb = props_a.get('virya', ''), props_b.get('virya', '') | |
| if label_idx == 2: return (f'Opposing Virya ({va} vs {vb}): classical Viruddha pattern.', 'Pitta/Vata vitiation, metabolic confusion.') | |
| if label_idx == 1: return (f'Pharmacological caution: Rasa/Virya overlap ({va}, {vb}).', 'Possible Pitta aggravation, dryness, or excess stimulation.') | |
| return ('Aligned Rasa/Virya supports combined therapeutic action.', 'Harmonized Agni and profound therapeutic effect.') | |
| def predict_pair(herb_a_display, herb_b_display): | |
| a, b = resolve_herb_name(herb_a_display), resolve_herb_name(herb_b_display) | |
| if a and b: | |
| key = tuple(sorted([a, b])) | |
| dict_key = f"{key[0]}|{key[1]}" | |
| if dict_key in CLASSICAL_INTERACTIONS: | |
| info = CLASSICAL_INTERACTIONS[dict_key] | |
| return {'type': info['type'], 'confidence': '100% (Classical Text)', 'source': 'Charaka Samhita Curated Pair', 'reason': info.get('reason', 'Documented classical interaction.'), 'body': info.get('body', 'Varies based on text.')} | |
| if not a or not b or interaction_model is None: | |
| return {'type': 'UNKNOWN', 'confidence': '0%', 'source': 'N/A', 'reason': 'One or both herbs missing from database.', 'body': 'N/A'} | |
| feat_a, feat_b = encode_herb(herb_props[a]), encode_herb(herb_props[b]) | |
| combined = np.concatenate([feat_a, feat_b, np.abs(feat_a - feat_b), feat_a * feat_b]) | |
| pred = int(interaction_model.predict([combined])[0]) | |
| proba = interaction_model.predict_proba([combined])[0] | |
| reason, body = rule_based_reason(herb_props[a], herb_props[b], pred) | |
| return {'type': LABEL_NAMES[pred], 'confidence': f"{proba[pred]*100:.1f}% (ML + Dravyaguna rules)", 'source': 'RandomForest on herb_properties', 'reason': reason, 'body': body} | |
| def generate_report(symptoms, season, age, gender): | |
| clinical_text = f"symptoms: {symptoms} | season: {season} | age: {age} | gender: {gender}" | |
| enc = tokenizer(clinical_text, return_tensors='pt', max_length=128, padding='max_length', truncation=True) | |
| with torch.no_grad(): out = model(enc['input_ids'].to(device), enc['attention_mask'].to(device)) | |
| herb_probs = torch.sigmoid(out['herb']).cpu().numpy()[0] | |
| top_idx = herb_probs.argsort()[::-1][:4] | |
| recommended = [(all_herbs[i].capitalize(), float(herb_probs[i])) for i in top_idx] | |
| dosha_pred = le_dosha.inverse_transform([torch.argmax(out['dosha'], 1).cpu().item()])[0] | |
| severity_pred = SEV_NAMES[torch.argmax(out['severity'], 1).cpu().item()] | |
| conflict = torch.sigmoid(out['conflict']).item() > 0.5 | |
| toxicity_pred = TOX_NAMES[torch.argmax(out['toxicity'], 1).cpu().item()] | |
| dosage_mg = float(out['dosage'].item() * 1000) | |
| se_probs = torch.sigmoid(out['side_effects']).cpu().numpy()[0] | |
| active_se = [side_effect_keywords[i] for i in range(len(se_probs)) if se_probs[i] > 0.5] if side_effect_keywords else [] | |
| lines = [] | |
| lines.append("=" * 70) | |
| lines.append(" AYURGENIX V9 — CLINICAL INTELLIGENCE REPORT") | |
| lines.append("=" * 70) | |
| lines.append(f"\nPatient: {age} yrs | {gender} | {season}") | |
| lines.append(f"Symptoms: {symptoms}") | |
| lines.append("\n--- MULTI-TASK ANALYSIS (Neural Network Predictions) ---") | |
| lines.append(f" Dosha imbalance : {dosha_pred}") | |
| lines.append(f" Disease severity : {severity_pred}") | |
| lines.append(f" Drug safety : {'CAUTION DETECTED' if conflict else 'No Conflicts'}") | |
| lines.append(f" Toxicity level : {toxicity_pred}") | |
| lines.append(f" Safe dosage (model) : {dosage_mg:.0f} mg/day") | |
| lines.append(f" Side effects : {', '.join(active_se) if active_se else 'None detected'}") | |
| lines.append("\n--- RECOMMENDED HERBS (Neural Network Predictions) ---") | |
| rec_canonical = [] | |
| for h, prob in recommended: | |
| canon = resolve_herb_name(h) or '?' | |
| rec_canonical.append(canon) | |
| preview = f"rasa={herb_props[canon].get('rasa','')} virya={herb_props[canon].get('virya','')}" if canon in herb_props else "" | |
| lines.append(f" -> {h} ({prob*100:.1f}%) [DB: {canon}]") | |
| if preview: lines.append(f" {preview}") | |
| lines.append("\n--- CLASSICAL FORMULATION MATCH (From Bhaishajya Kosha) ---") | |
| best, best_score = None, -1 | |
| sym_tokens = tokenize_indications(symptoms) | |
| herb_set = set([c for c in rec_canonical if c]) | |
| for form in formulations: | |
| ftokens = tokenize_indications(form.get('indications', '')) | |
| overlap = len(sym_tokens & ftokens) * 2 + len(herb_set & set([resolve_herb_name(i) for i in form.get('main_ingredients', []) if resolve_herb_name(i)])) * 3 | |
| if overlap > best_score: best_score, best = overlap, form | |
| if best: | |
| lines.append(f" Name : {best.get('name', 'N/A')}") | |
| lines.append(f" Category : {best.get('category', 'N/A')}") | |
| lines.append(f" Dosage : {best.get('dosage', 'Standard Dosage')}") | |
| lines.append(f" Anupana : {best.get('anupana', 'Warm water')}") | |
| lines.append(f" Reference : {best.get('reference', 'Classical Text')}") | |
| lines.append(f" Indications: {best.get('indications', '')[:200]}...") | |
| else: lines.append(" No direct classical formulation match found.") | |
| lines.append("\n--- HERB-HERB INTERACTIONS (From KnowledgeBundle & ML) ---") | |
| names = [h for h, _ in recommended] | |
| interactions_found = 0 | |
| for i in range(len(names)): | |
| for j in range(i + 1, len(names)): | |
| pair_res = predict_pair(names[i], names[j]) | |
| if pair_res['type'] == 'UNKNOWN': continue | |
| interactions_found += 1 | |
| tag = {'SYNERGISTIC': '[OK]', 'CAUTION': '[WARN]', 'CONTRAINDICATED': '[DANGER]'}.get(pair_res['type'], '[?]') | |
| lines.append(f"\n {tag} {names[i]} + {names[j]}") | |
| lines.append(f" Type : {pair_res['type']} | {pair_res['confidence']}") | |
| lines.append(f" Source : {pair_res['source']}") | |
| lines.append(f" Reason : {pair_res['reason']}") | |
| lines.append(f" Body effect: {pair_res['body']}") | |
| if interactions_found == 0: lines.append(" No documented interactions found in database for the recommended herbs.") | |
| lines.append("\n--- DISCLAIMER ---") | |
| lines.append("AI-assisted Ayurvedic clinical decision support for research and education only.") | |
| lines.append("=" * 70) | |
| return "\n".join(lines) | |
| interface = gr.Interface( | |
| fn=generate_report, | |
| inputs=[gr.Textbox(label="Symptoms"), gr.Dropdown(["Summer", "Winter", "Monsoon", "Autumn", "Spring"], label="Season"), gr.Number(label="Age", value=30), gr.Dropdown(["Male", "Female"], label="Gender")], | |
| outputs=gr.Textbox(label="Clinical Report") | |
| ) | |
| interface.launch() | |