import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
from huggingface_hub import login

# Login to Hugging Face using the token
read_key = os.environ.get("Read_Token")
login(token=read_key)

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b-it",
    torch_dtype=torch.bfloat16
)

# Streamlit UI
st.title("Gemma-2B-IT Text Generator")
st.write("Enter the text and generate a structured prompt.")

# User inputs the text
text = st.text_area("Input your text", 
    "El sistema CRISPR-Cas9 permite una edición precisa del genoma mediante la creación de rupturas de doble cadena en ubicaciones específicas del ADN, lo que facilita modificaciones genéticas específicas.")

# Create the base prompt structure
base_prompt = f"""
Realiza las siguientes acciones:
1 - Tradúceme el texto delimitado por las comillas triples a los siguientes 3 idiomas: inglés, árabe y francés.
2 - Identifícame el tema principal del texto delimitado por las comillas triples.
3 - Obtén el tono en el que está escrito el texto delimitado por las comillas triples.
4 - Devuelve el resultado de todos los pasos en 1 solo JSON con los siguientes headers:TraduccionIngles, TraduccionArabe, TraduccionFrances, TemaPrincipal, Tono

Text:
```{text}```
"""

# Allow the user to modify the prompt dynamically
custom_prompt = st.text_area("Edit Prompt Template", base_prompt, height=300)

# Sliders for generation parameters
max_length = st.slider("Max Length", min_value=50, max_value=500, value=400)
temperature = st.slider("Temperature", min_value=0.1, max_value=1.5, value=0.7)
top_p = st.slider("Top-p (nucleus sampling)", min_value=0.0, max_value=1.0, value=0.9)
repetition_penalty = st.slider("Repetition Penalty", min_value=1.0, max_value=2.0, value=1.1)
do_sample = st.checkbox("Enable Sampling", value=True)

# Generate text when the button is pressed
if st.button("Generate"):
    # Correctly format the prompt using the current text input
    formatted_prompt = custom_prompt.format(text=text)

    input_ids = tokenizer(formatted_prompt, return_tensors="pt")
    
    # Generate text with specified parameters
    outputs = model.generate(
        **input_ids, 
        max_length=max_length, 
        num_return_sequences=1, 
        do_sample=do_sample, 
        temperature=temperature, 
        top_p=top_p, 
        repetition_penalty=repetition_penalty
    )
    
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Display the generated text
    st.write("Generated Text:")
    st.text_area("Output", generated_text, height=300)