Update app.py

app.py

@@ -1,3 +1,26 @@
-import gradio as gr
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-gr.Interface.load("models/AkshatSurolia/ICD-10-Code-Prediction").launch()
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("AkshatSurolia/ICD-10-Code-Prediction")
+model = AutoModelForSequenceClassification.from_pretrained("AkshatSurolia/ICD-10-Code-Prediction")
+
+# Create a Streamlit input text box
+input_text = st.text_input("Enter your text:")
+
+# If input is provided
+if input_text:
+    # Limit the input length
+    truncated_input = input_text[:512]
+    
+    # Tokenize the input
+    tokens = tokenizer(truncated_input, truncation=True, padding=True, return_tensors="pt")
+
+    # Get model output
+    output = model(**tokens)
+
+    # The output of the model is a logits vector, so we take the argmax to get the predicted class index
+    predicted_class_idx = torch.argmax(output.logits, dim=-1).item()
+    
+    st.write(f"Predicted class index: {predicted_class_idx}")