app.py

import torch
import spaces
import gradio as gr

# from transformers import (
#     AutoModelForCausalLM,
#     AutoTokenizer,
#     AutoModelForSequenceClassification,
# )

# chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
# chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, dtype=torch.bfloat16, device_map="cpu")
# chat_model.to("cuda")
# chat_tokenizer = AutoTokenizer.from_pretrained(chat_model_name)

# moderator_model_name = "saiteki-kai/QA-DeBERTa-v3-large-binary-3"
# moderator_model = AutoModelForSequenceClassification.from_pretrained(moderator_model_name, device_map="cpu")
# moderator_model.to("cuda")
# moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name, padding_side="right")

# def generate_responses(model, tokenizer, prompts):
#     messages = [[{"role": "user", "content": message}] for message in prompts]

#     texts = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
#     model_inputs = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(model.device)

#     with torch.inference_mode():
#         generated_ids = model.generate(
#             **model_inputs,
#             do_sample=False,
#             temperature=0,
#             repetition_penalty=1.1,
#             max_new_tokens=512,
#         )
#     prompt_lengths = model_inputs["attention_mask"].sum(dim=1) - 1
#     generated_ids = [
#         output_ids[length:] for length, output_ids in zip(prompt_lengths, generated_ids)
#     ]
#     responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

#     return responses

# def classify_pairs(model, tokenizer, prompts, responses):
#     texts = [
#         prompt + "[SEP]" + response for prompt, response in zip(prompts, responses)
#     ]

#     input_ids = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(model.device)
#     print(tokenizer.batch_decode(input_ids["input_ids"]))

#     with torch.inference_mode():
#         outputs = model(**input_ids)
#         scores = torch.softmax(outputs.logits, dim=-1).detach().cpu()
#         unsafety_scores = [float(s[1]) for s in scores] # get unsafe axis

#     return unsafety_scores


@spaces.GPU(duration=120)
def generate(submission: list[dict[str, str]], team_id: str) -> list[dict[str, str | float]]:
    print("GENERATE")

    # ids = [s["id"] for s in submission]
    # prompts = [s["prompt"] for s in submission]

    # responses = generate_responses(chat_model, chat_tokenizer, prompts)
    # print(responses)

    # scores = classify_pairs(moderator_model, moderator_tokenizer, prompts, responses)
    # print(scores)

    chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
    ids = [s["id"] for s in submission]
    prompts = [s["prompt"] for s in submission]
    responses = ["This is a placeholder response." for _ in prompts]
    scores = [0.5 for _ in prompts]

    outputs = [
        {"id": id, "prompt": prompt, "response": response, "score": score, "model": chat_model_name, "team_id": team_id}
        for id, prompt, response, score in zip(ids, prompts, responses, scores)
    ]

    return outputs


with gr.Blocks() as demo:
    print("START")
    gr.api(generate, api_name="scores", concurrency_limit=None, batch=False)


print("LAUNCH")
demo.launch()