app.py

app.py

@@ -1,38 +1,20 @@
-import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-def predict_rooms(new_students, new_temperature):
-    # Load the model and tokenizer
-    model_name = "AI"  # Replace with the name or path of the model you want to use
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=1)
 
-    # Convert the input to tokens
-    inputs = tokenizer.encode_plus(
-        "Number of students: {}, Temperature: {}".format(new_students, new_temperature),
-        padding="max_length",
-        truncation=True,
-        max_length=64,
-        return_tensors="pt"
-    )
+# Input text
+text = "I want to book a flight from New York to London on July 1st."
 
-    # Make the prediction
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        predicted_rooms = torch.argmax(logits, dim=1).item()
+# Tokenize input text
+inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
 
-    return predicted_rooms
+# Predict flight ticket
+with torch.no_grad():
+    logits = model(**inputs).logits
+    prediction = torch.sigmoid(logits).item()
 
-def greet(name):
-    return "Hello " + name + "!"
-
-iface = gr.Interface(
-    fn=[predict_rooms, greet],
-    inputs=[["number", "number"], "text"],
-    outputs=["number", "text"],
-    title="Room Prediction",
-    description="Predict the number of rooms based on the number of students and temperature, and greet the user."
-)
-
-iface.launch()
+# Display prediction
+print("Flight ticket prediction:", prediction)