BioMistral-7B-SymptomDiagnosis

from transformers import AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig
from peft import PeftModel
import torch

# Load model
MODEL_NAME = "Sugandha-Chauhan/BioMistral-7B-SymptomDiagnosis"
BASE_MODEL = "BioMistral/BioMistral-7B"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForSequenceClassification.from_pretrained(
    BASE_MODEL,
    num_labels=10,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.float16
)

model = PeftModel.from_pretrained(model, MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Predict
def predict(symptoms_text):
    inputs = tokenizer(
        symptoms_text,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128
    )
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    probabilities = torch.softmax(logits, dim=-1)
    confidence, predicted_class = torch.max(probabilities, dim=-1)
    
    diagnosis_classes = {
        0: "acute bronchitis",
        1: "anxiety",
        2: "conjunctivitis due to allergy",
        3: "eczema",
        4: "infectious gastroenteritis",
        5: "pneumonia",
        6: "psoriasis",
        7: "spondylosis",
        8: "sprain or strain",
        9: "strep throat"
    }
    
    return diagnosis_classes[predicted_class.item()], confidence.item()

# Example
diagnosis, confidence = predict("nausea ,vomiting ,diarrhea and fever")
print(f"Diagnosis: {diagnosis} ({confidence:.1%} confidence)")