import gradio as gr
from transformers import pipeline
import torch
import os

hf_token = os.environ.get("HF_TOKEN")

# Load model TxGemma-2B
print("Load model TxGemma...")
pipe = pipeline(
    "text-generation",
    model="google/txgemma-2b-predict",
    device="cpu",
    torch_dtype=torch.float32,
    token=hf_token 
)

def process_prompt(user_input):
    """
    Processing
    """
    try:
        # Generate response
        outputs = pipe(
            user_input, 
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7
        )
        response = outputs[0]["generated_text"]
        return response
    except Exception as e:
        return f"Errore durante l'elaborazione: {str(e)}"

# Create Interface
demo = gr.Interface(
    fn=process_prompt,
    inputs=gr.Textbox(
        label="Inserisci il prompt per TxGemma",
        placeholder="Esempio: CCO (molecola di etanolo)",
        lines=5
    ),
    outputs=gr.Textbox(
        label="Risposta di TxGemma",
        lines=10
    ),
    title="Demo TxGemma - Analisi Molecolare",
    description="""
            Inserisci un prompt (es. stringa SMILES di una molecola) e TxGemma fornirà previsioni sulle proprietà terapeutiche.
            
            ---
            <br><br>
            **Created with ❤️ by Rocco for Giulio(GOD)**
        """,
    theme="soft"
)

# Launch app
if __name__ == "__main__":
    demo.launch()

# After load of the model
model_info = f"""
**Model load:** {pipe.model.config._name_or_path}  
**Type:** {pipe.model.config.model_type}  
**Parameters:** {pipe.model.num_parameters() / 1e9:.2f}B  
**Spec:** Therapeutic Development (TDC)
"""