Update app.py

app.py

@@ -1,24 +1,38 @@
-import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Obter o token de autenticação a partir do Secret
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+checkpoint = "jinaai/reader-lm-0.5b"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Definir o modelo e o tokenizador
-model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HUGGINGFACE_TOKEN, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=HUGGINGFACE_TOKEN, trust_remote_code=True)
+# Carrega modelo e tokenizer
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=200)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
+def process_html(html_content):
+    messages = [{"role": "user", "content": html_content}]
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False)
+    
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(
+        inputs,
+        max_new_tokens=1024,
+        temperature=0,
+        do_sample=False,
+        repetition_penalty=1.08
+    )
+    
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return decoded_output
 
-# Criar a interface Gradio
-iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", live=True, title="Geração de Texto com Mixtral")
+# Interface Gradio
+iface = gr.Interface(
+    fn=process_html,
+    inputs=gr.Textbox(lines=10, placeholder="Insira o conteúdo HTML aqui...", label="HTML"),
+    outputs=gr.Textbox(label="Resposta do modelo"),
+    title="HTML Reader (jinaai/reader-lm-0.5b)",
+    description="Insira um conteúdo HTML e veja como o modelo interpreta o conteúdo."
+)
 
-# Executar o app
 if __name__ == "__main__":
     iface.launch()