Update app.py

app.py

@@ -1,15 +1,23 @@
+import os
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
+import re
+
+# Securely get the Hugging Face token
+hf_token = os.environ.get("HUGGINGFACE_TOKEN")
+if not hf_token:
+    raise ValueError("HUGGINGFACE_TOKEN not found in environment variables")
 
 # Load the model and tokenizer
 base_model_name = "meta-llama/Llama-3.2-1B"
 fine_tuned_model_path = "./gross_llama"  # Path to your fine-tuned model
 
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
-peft_model = PeftModel.from_pretrained(base_model, fine_tuned_model_path)
+# Use the token for authentication
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, token=hf_token)
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, token=hf_token)
+peft_model = PeftModel.from_pretrained(base_model, fine_tuned_model_path, token=hf_token)
 
 peft_model.to("cuda")
 peft_model.eval()
@@ -25,27 +33,18 @@ Sign Language Translator for converting sentences to Gloss. Glosses are written
 ### Response:
 """
 
-import re
-
 def translate_to_gloss(input_text):
-    inputs = tokenizer(
-        [tinyllama_prompt.format(input_text)],
-        return_tensors="pt"
-    ).to("cuda")
-
+    inputs = tokenizer([tinyllama_prompt.format(input_text)], return_tensors="pt").to("cuda")
     with torch.no_grad():
         outputs = peft_model.generate(**inputs, max_new_tokens=64, use_cache=True)
-    
     decoded_output = tokenizer.batch_decode(outputs)[0]
     
     # Extract the response part
     response_start = decoded_output.find("### Response")
     if response_start != -1:
         response = decoded_output[response_start + len("### Response"):].strip()
-        
         # Remove any remaining prompt parts
         response = re.sub(r'###.*$', '', response, flags=re.DOTALL).strip()
-        
         # Remove any non-gloss text (assuming gloss is in all caps)
         gloss_parts = re.findall(r'\b[A-Z]+(?:\s+[A-Z]+)*\b', response)
         gloss = ' '.join(gloss_parts)