app.py

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

device = torch.device('cpu') # HF Free Tier uses CPU
ckpt = torch.load("AyurGenixV9_FULLY_EMBEDDED.pt", 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()