Update app.py

app.py

@@ -3,9 +3,12 @@
 import gradio as gr
 import os
 
+import torch
 from peft import PeftConfig, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 base_model_name = "google/gemma-7b"
 #adapter_model_name = "samidh/cope-g2b-2c-hs-skr-s1.5.9-sx-sk-s5.d25"
 #adapter_model_name = "samidh/cope-g2b-2c-hs-skr-s1.5.9-sx-sk-s1.5.l1e4-e10-d25"
@@ -20,6 +23,8 @@ model = AutoModelForCausalLM.from_pretrained(base_model_name, token=os.environ['
 model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
 model.merge_and_unload()
 
+model = model.to(device)
+
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
 #inputs = tokenizer.encode("This movie was really", return_tensors="pt")
@@ -90,14 +95,17 @@ DEFAULT_CONTENT = "LLMs steal our jobs."
 # Function to make predictions
 def predict(content, policy):
     input_text = PROMPT.format(policy=policy, content=content)
-    inputs = tokenizer.encode(input_text, return_tensors="pt")
-    outputs = model.generate(inputs, max_new_tokens=1)
-    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    if int(decoded_output[-1]) == 0:
-        return f'NON-Violating ({decoded_output[-1]})'
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model(input_ids)
+        logits = outputs.logits[:, -1, :]  # Get logits for the last token
+        predicted_token_id = torch.argmax(logits, dim=-1).item()
+    decoded_output = tokenizer.decode([predicted_token_id])
+    if decoded_output == '0':
+        return f'NON-Violating ({decoded_output})'
     else:
-        return f'VIOLATING ({decoded_output[-1]})'
-
+        return f'VIOLATING ({decoded_output})'
 
 with gr.Blocks() as iface:
     gr.Markdown("# CoPE Alpha Preview")