feat: add logging

app.py

@@ -1,5 +1,6 @@
 import torch
 import spaces
+import logging
 import gradio as gr
 
 from transformers import (
@@ -8,13 +9,18 @@ from transformers import (
     AutoModelForSequenceClassification,
 )
 
+logging.basicConfig(level=logging.INFO)
+
+logger = logging.getLogger(__name__)
+
+
 chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
-chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, dtype=torch.bfloat16)
+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"
-moderator_model = AutoModelForSequenceClassification.from_pretrained(moderator_model_name)
+moderator_model = AutoModelForSequenceClassification.from_pretrained(moderator_model_name, device_map="cpu")
 moderator_model.to("cuda")
 moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name)
 
@@ -65,11 +71,15 @@ def generate(submission: list[dict[str, str]], team_id: str) -> list[dict[str, s
     responses = generate_responses(chat_model, chat_tokenizer, prompts)
     scores = classify_pairs(moderator_model, moderator_tokenizer, prompts, responses)
     
-    return [
+    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)
     ]
 
+    print(outputs)
+    
+    return outputs
+
 
 with gr.Blocks() as demo:
     gr.Markdown("Welcome")