Update app.py

app.py

@@ -2,35 +2,22 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-from huggingface_hub import login
-
-from huggingface_hub import login
-import os
-
-hf_token = os.getenv("HF_TOKEN")  # Fetch token securely
-login(hf_token)
-# Load Mixtral model
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-# Define the model name
-model_name = "microsoft/phi-2"  # Ensure this is correct
-
-# Load the tokenizer and model
+# Load the model and tokenizer
+model_name = "microsoft/phi-2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Set the device (use GPU if available, otherwise fallback to CPU)
+# Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)  # Move model to the appropriate device
-
-print("Model loaded successfully!")
+model.to(device)
 
-# Create Gradio UI
-iface = gr.Interface(fn=solve_problem, 
-                     inputs=gr.Textbox(lines=2, placeholder="Enter a JEE/NEET question..."), 
-                     outputs="text",
-                     title="Prof. Cool - JEE/NEET Solver",
-                     description="Ask a complex JEE Advanced or NEET question and get a detailed solution.")
+# Define the function to process user input
+def solve_problem(user_input):
+    inputs = tokenizer(user_input, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_length=200)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
+# Create Gradio interface
+iface = gr.Interface(fn=solve_problem, inputs="text", outputs="text")
 iface.launch()