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

# 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.eval()

tinyllama_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
Sign Language Translator for converting sentences to Gloss. Glosses are written representations of sign language that capture the core meaning and grammatical structure of signed utterances.

### Input:
{}

### Response:
"""

def translate_to_gloss(input_text):
    inputs = tokenizer([tinyllama_prompt.format(input_text)], return_tensors="pt")
    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)
    else:
        gloss = "Error: Couldn't find the response in the output."
    
    return gloss

# Update the Gradio interface
iface = gr.Interface(
    fn=translate_to_gloss,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence to translate to gloss..."),
    outputs=gr.Textbox(label="Gloss Translation"),
    title="Sign Language Gloss Translator",
    description="This app translates English sentences into sign language glosses.",
    examples=[
        ["What are you doing today?"],
        ["I'm going to the park."],
        ["How are you feeling?"],
    ]
)

# Launch the app
iface.launch()