Description, model inference

app.py

@@ -1,7 +1,7 @@
 import spaces
 import gradio as gr
 import torch
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import json
 from datetime import datetime
 from uuid import uuid4
@@ -13,7 +13,6 @@ from huggingface_hub import CommitScheduler
 # TODO make it so that feedback is only saved on prev. example if user makes another obfuscation
 # and changes slider but doesn't hit obfuscate
 
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
 MODEL_PATHS = {
     "length_more": "hallisky/lora-length-long-llama-3-8b",
     "length_less": "hallisky/lora-length-long-llama-3-8b",
@@ -23,6 +22,24 @@ MODEL_PATHS = {
     "grade_less": "hallisky/lora-grade-elementary-llama-3-8b",
 }
 
+DESCRIPTION = """\
+# Authorship Obfuscation
+This Space demonstrates model [Llama-2-7b-chat](https://huggingface.co/meta-llama/Llama-2-7b-chat) by Meta, a Llama 2 model with 7B parameters fine-tuned for chat instructions. Feel free to play with it, or duplicate to run generations without a queue! If you want to run your own service, you can also [deploy the model on Inference Endpoints](https://huggingface.co/inference-endpoints).
+🔎 For more details about the Llama 2 family of models and how to use them with `transformers`, take a look [at our blog post](https://huggingface.co/blog/llama2).
+🔨 Looking for an even more powerful model? Check out the [13B version](https://huggingface.co/spaces/huggingface-projects/llama-2-13b-chat) or the large [70B model demo](https://huggingface.co/spaces/ysharma/Explore_llamav2_with_TGI).
+"""
+
+# Load models
+if not torch.cuda.is_available():
+    device = "cpu"
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+
+if torch.cuda.is_available():
+    device = "cuda"
+    model_id = "meta-llama/Meta-Llama-3-8B"
+    model = AutoModelForCausalLM.from_pretrained(model_id).to(device) # device_map="auto" requires accelerate
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
 # Global variable to store the latest obfuscation result
 latest_obfuscation = {}
 user_id = str(uuid4())  # Generate a unique session-specific user ID
@@ -40,11 +57,6 @@ scheduler = CommitScheduler(
     every=0.5
 )
 
-@spaces.GPU
-def temp(text):
-    response = tokenizer(text, return_tensors="pt")
-    return response
-
 def save_data(data):
     with scheduler.lock:
         with JSON_DATASET_PATH.open("a") as f:
@@ -55,13 +67,12 @@ def save_feedback(feedback_rating, feedback_text):
     global latest_obfuscation
     if latest_obfuscation:
         feedback_data = latest_obfuscation.copy()
-        feedback_data['feedback'] = {
-            "rating": feedback_rating,
-            "text": feedback_text
-        }
+        feedback_data["feedback_rating"] = feedback_rating
+        feedback_data["feedback_text"] = feedback_text
         save_data(feedback_data)
     return "No Feedback Selected", ""
 
+@spaces.GPU
 def greet(input_text, length, function_words, grade_level, sarcasm, formality, voice, persuasive, descriptive, narrative, expository):
     global latest_obfuscation, user_id
     current_time = datetime.now().isoformat()
@@ -80,7 +91,15 @@ def greet(input_text, length, function_words, grade_level, sarcasm, formality, v
         f"Narrative: {narrative}\n"
         f"Expository: {expository}"
     )
-    
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=tokenizer(input_text, return_tensors="pt").input_ids.to(device),
+            max_length=100,
+            num_return_sequences=1,
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
     # Save the new obfuscation result and reset feedback
     latest_obfuscation = {
         "datetime": current_time,
@@ -99,10 +118,8 @@ def greet(input_text, length, function_words, grade_level, sarcasm, formality, v
             "expository": expository
         },
         "output": response,
-        "feedback": {
-            "rating": "No Feedback Selected",
-            "text": ""
-        }
+        "feedback_rating": "No Feedback Selected",
+        "feedback_text": ""
     }
     
     # Save the obfuscation result
@@ -144,6 +161,7 @@ def check_initial_feedback_state(feedback_rating, feedback_text):
 demo = gr.Blocks()
 
 with demo:
+    gr.Markdown(DESCRIPTION)
     with gr.Row():
         with gr.Column(variant="panel"):
             gr.Markdown("# 1) Input Text\n### Enter the text to be obfuscated.")
@@ -206,6 +224,9 @@ with demo:
             # Initialize the button and warning message state on page load
             demo.load(fn=update_obfuscate_button, inputs=input_text, outputs=[obfuscate_button, warning_message])
 
+        # with gr.Column(variant="panel"):
+        #     gr.Markdown("# 3) Obfuscated Output")
+
         with gr.Column(variant="panel"):
             gr.Markdown("# 3) Obfuscated Output")