Update app.py

app.py

@@ -1,76 +1,40 @@
 import torch
 import torch.nn as nn
-import gradio as gr
 from transformers import AutoTokenizer, Qwen2ForCausalLM
 
 device = "cpu"
 dtype = torch.float32
 
-# Qwen 0.5B = hidden_size 896
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
 text_encoder = Qwen2ForCausalLM.from_pretrained(
     "Qwen/Qwen2-0.5B",
     torch_dtype=dtype,
 )
 
-# Projection 896 -> 2048 pour FLUX.1-Schnell
-proj = nn.Linear(896, 2048)
+# Qwen 0.5B → 896 dims
+proj_tokens = nn.Linear(896, 2048)   # pour prompt_embeds
+proj_pooled = nn.Linear(2048, 768)   # pour pooled_prompt_embeds
 
-def encode(prompt: str):
-    # Nettoyage du prompt
-    if prompt is None:
-        prompt = ""
-    prompt_clean = prompt.strip()
+def encode(prompt):
+    if not prompt.strip():
+        prompt = tokenizer.eos_token or "."
 
-    # Si vide -> token de secours
-    if prompt_clean == "":
-        if tokenizer.eos_token:
-            prompt_clean = tokenizer.eos_token
-        else:
-            prompt_clean = "."
+    tokens = tokenizer(prompt, return_tensors="pt")
 
-    # Tokenisation
-    tokens = tokenizer(
-        prompt_clean,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=512,
-    )
-
-    # Encodage Qwen (SANS inference_mode)
     out = text_encoder(
         **tokens,
         output_hidden_states=True,
         use_cache=False,
     )
 
-    # Embeddings Qwen (896 dims)
-    embeds_896 = out.hidden_states[-1]  # [1, L, 896]
-
-    # Projection -> 2048 dims
-    embeds_2048 = proj(embeds_896)      # [1, L, 2048]
+    hidden = out.hidden_states[-1]  # [1, L, 896]
 
-    # pooled -> moyenne
-         # [1, 2048]
-    proj_pooled = nn.Linear(2048, 768)
-    pooled = proj_pooled(embeds_2048.mean(dim=1))
-    
-    # Sauvegarde
-    torch.save(embeds_2048, "embeds.pt")
-    torch.save(pooled, "pooled.pt")
+    embeds_2048 = proj_tokens(hidden)  # [1, L, 2048]
 
-    return str(embeds_2048.shape), "embeds.pt", "pooled.pt"
+    pooled_2048 = embeds_2048.mean(dim=1)  # [1, 2048]
+    pooled_768 = proj_pooled(pooled_2048)  # [1, 768]
 
-demo = gr.Interface(
-    fn=encode,
-    inputs=gr.Textbox(label="Prompt"),
-    outputs=[
-        gr.Textbox(label="Shape"),
-        gr.File(label="Embeddings 2048"),
-        gr.File(label="Pooled 2048")
-    ],
-    title="External Text Encoder — 2048 dims (FLUX.1‑Schnell)"
-)
+    torch.save(embeds_2048, "embeds.pt")
+    torch.save(pooled_768, "pooled.pt")
 
-demo.launch()
+    return "OK", "embeds.pt", "pooled.pt"