from fastai.text.all import *
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import gradio as gr

# Load the medical model
medical_learn = load_learner('model.pkl')

# Medical model configuration
medical_description = "Medical Diagnosis"
medical_categories = ['Allergy', 'Anemia', 'Bronchitis', 'Diabetes', 'Diarrhea', 'Fatigue', 'Flu', 'Malaria', 'Stress']

def classify_medical_text(txt):
    pred, idx, probs = medical_learn.predict(txt)
    return dict(zip(medical_categories, map(float, probs)))

# Load the psychiatric model from Hugging Face
psychiatric_model_name = "mental/mental-bert-base-uncased"  # Replace with the appropriate model
psychiatric_tokenizer = AutoTokenizer.from_pretrained(psychiatric_model_name)
psychiatric_model = AutoModelForSequenceClassification.from_pretrained(psychiatric_model_name)

# Psychiatric model configuration
psychiatric_description = "Psychiatric Analysis"
psychiatric_labels = ['Depression', 'Anxiety', 'Bipolar Disorder', 'PTSD', 'OCD', 'Stress', 'Schizophrenia']  # Adjust based on the model

def classify_psychiatric_text(txt):
    inputs = psychiatric_tokenizer(txt, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = psychiatric_model(**inputs)
    logits = outputs.logits
    probabilities = torch.softmax(logits, dim=1).squeeze().tolist()
    return dict(zip(psychiatric_labels, probabilities))

# Gradio Interfaces
medical_text = gr.Textbox(lines=2, label='Describe your symptoms in detail')
medical_label = gr.Label()
medical_examples = ['I feel short of breath and have a high fever.', 'My throat hurts and I keep sneezing.', 'I am always thirsty.']

psychiatric_text = gr.Textbox(lines=2, label='Describe your mental health concerns in detail')
psychiatric_label = gr.Label()
psychiatric_examples = ['I feel hopeless and have no energy.', 'I am unable to concentrate and feel anxious all the time.', 'I have recurring intrusive thoughts.']

medical_interface = gr.Interface(
    fn=classify_medical_text,
    inputs=medical_text,
    outputs=medical_label,
    examples=medical_examples,
    description=medical_description,
)

psychiatric_interface = gr.Interface(
    fn=classify_psychiatric_text,
    inputs=psychiatric_text,
    outputs=psychiatric_label,
    examples=psychiatric_examples,
    description=psychiatric_description,
)

# Combine interfaces using Tabs
app = gr.TabbedInterface([medical_interface, psychiatric_interface], ["Medical Diagnosis", "Psychiatric Analysis"])
app.launch(inline=False)