File size: 2,649 Bytes
9228688
81f65fc
9228688
28f7004
81f86f8
9228688
 
f86c935
 
5220ce7
 
28f7004
5220ce7
28f7004
5220ce7
28f7004
 
 
 
 
 
 
 
 
5220ce7
f86c935
28f7004
f86c935
5220ce7
f86c935
81f86f8
f86c935
 
 
 
 
 
 
 
 
 
 
9228688
 
f86c935
 
 
 
 
81f86f8
f86c935
 
 
81f86f8
f86c935
 
 
 
 
 
 
 
 
 
 
 
5220ce7
f86c935
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import streamlit as st
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import os

@st.cache_resource
def load_model():
    """Load your trained model"""
    try:
        # Load from the same directory as the script
        model_path = os.path.dirname(__file__)
        
        st.info(f"Loading model from: {model_path}")
        
        # Load model
        model = DistilBertForSequenceClassification.from_pretrained(
            model_path,
            local_files_only=True
        )
        tokenizer = DistilBertTokenizer.from_pretrained(
            model_path,
            local_files_only=True
        )
        
        st.success("✅ Model loaded successfully!")
        return model, tokenizer
        
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

def predict_text(text, model, tokenizer):
    """Make prediction"""
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        truncation=True, 
        padding=True,
        max_length=128
    )
    
    model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(predictions, dim=-1).item()
        confidence = predictions[0][predicted_class].item()
    
    return predicted_class, confidence

# Main App
st.title("🏥 Medical Text Classifier")
st.write("Enter text to classify as Medical or Non-Medical")

# Load model
model, tokenizer = load_model()

if model is not None:
    # Text input
    user_input = st.text_area(
        "Enter your text:",
        placeholder="Example: I have a headache and need to see a doctor...",
        height=100
    )
    
    # Classify button
    if st.button("🔍 Classify Text", type="primary"):
        if user_input.strip():
            with st.spinner("Analyzing..."):
                predicted_class, confidence = predict_text(user_input, model, tokenizer)
                
                # Show results
                labels = ["Non-Medical", "Medical"]
                result = labels[predicted_class]
                
                if predicted_class == 1:  # Medical
                    st.success(f"🏥 **{result}**")
                else:  # Non-Medical
                    st.info(f"ℹ️ **{result}**")
                
                st.write(f"**Confidence:** {confidence:.1%}")
                st.progress(confidence)
        else:
            st.warning("Please enter some text to classify!")

else:
    st.error("Failed to load the model. Please check the files.")