import gradio as gr
import json
from sentence_transformers import SentenceTransformer, util

# Load the model
model = SentenceTransformer('BAAI/bge-small-en-v1.5')

def grade_logic(student_json, correct_json):
    try:
        s_list = json.loads(student_json)
        c_list = json.loads(correct_json)
        
        s_embs = model.encode(s_list, normalize_embeddings=True, convert_to_tensor=True)
        c_embs = model.encode(c_list, normalize_embeddings=True, convert_to_tensor=True)
        
        scores = []
        for i in range(len(s_list)):
            ans = str(s_list[i]).strip()
            if not ans:
                scores.append(0)
                continue
            sim = util.cos_sim(s_embs[i], c_embs[i]).item()
            if sim >= 0.80: 
                scores.append(100)
            elif sim <= 0.40: 
                scores.append(0)
            else: 
                scores.append(int(((sim - 0.40) / (0.80 - 0.40)) * 100))
        return json.dumps(scores)
    except Exception as e:
        return json.dumps([-1, str(e)])

# Create the interface and disable the queue here instead of in launch()
demo = gr.Interface(
    fn=grade_logic, 
    inputs=["text", "text"], 
    outputs="text",
    api_name="grade"
)

# Launch with only valid Gradio 4 parameters
demo.launch(share=False)