from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
import tensorflow as tf
import gradio as gr

# Load the tokenizer and model
model_name = "Zabihin/Symptom_to_Diagnosis"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = TFAutoModelForSequenceClassification.from_pretrained(model_name)

# Clean the input text
def clean_input(symptom_text):
    # Remove unwanted characters or non-ASCII characters
    symptom_text = ''.join([c for c in symptom_text if ord(c) < 128])
    symptom_text = symptom_text.lower()  # Optional: Convert to lowercase
    return symptom_text

# Define the predict function
def predict(symptom_text, chat_history=[]):
    try:
        # Clean the input
        symptom_text = clean_input(symptom_text)

        # Tokenize the input
        inputs = tokenizer(symptom_text, return_tensors="tf", padding=True, truncation=True, max_length=512)

        # Get model output
        outputs = model(**inputs)
        logits = outputs.logits
        prediction = tf.argmax(logits, axis=-1).numpy()[0]

        # Map the prediction to a label
        labels = {
            0: "Allergy", 1: "Arthritis", 2: "Bronchial Asthma", 3: "Cervical Spondylosis", 
            4: "Chicken Pox", 5: "Common Cold", 6: "Dengue", 7: "Diabetes", 8: "Drug Reaction", 
            9: "Fungal Infection", 10: "Gastroesophageal Reflux Disease", 11: "Hypertension", 
            12: "Impetigo", 13: "Jaundice", 14: "Malaria", 15: "Migraine", 16: "Peptic Ulcer Disease", 
            17: "Pneumonia", 18: "Psoriasis", 19: "Typhoid", 20: "Urinary Tract Infection", 21: "Varicose Veins"
        }

        diagnosis = labels.get(prediction, "Unknown diagnosis")

        # Add conversation history
        chat_history.append(("User", symptom_text))
        chat_history.append(("AI", f"Predicted Diagnosis: {diagnosis}. Please consult a doctor for more accurate results."))

    except Exception as e:
        chat_history.append(("AI", f"Error: {str(e)}"))

    return chat_history, ""

# Gradio UI
with gr.Blocks() as interface:
    gr.Markdown("<h1 style='text-align: center; font-size: 40px; margin-top: 30px; margin-bottom: 30px;'>Medi Mind - Your AI Health Assistant</h1>")
    chatbot = gr.Chatbot()
    input_box = gr.Textbox(show_label=False, placeholder="Describe your symptoms here...")
    send_button = gr.Button("Send")

    input_box.submit(predict, [input_box, chatbot], [chatbot, input_box])
    send_button.click(predict, [input_box, chatbot], [chatbot, input_box])

if __name__ == "__main__":
    interface.launch(share=True, server_name="0.0.0.0", server_port=7860, debug=True)
h