Update app.py

app.py

@@ -1,13 +1,9 @@
 import streamlit as st
-from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
 import PyPDF2
 import pandas as pd
 import torch
-import torch.nn as nn  # Added this import
-import numpy as np
-from copy import deepcopy
-import math
 import time
 
 # Device setup
@@ -21,7 +17,7 @@ st.set_page_config(
 )
 
 # Model name
-MODEL_NAME = "deepseek-ai/DeepSeek-V3-0324" #"amiguel/en2fr-transformer"
+MODEL_NAME = "deepseek-ai/DeepSeek-V3-0324"
 
 # Translation prompt template
 TRANSLATION_PROMPT = """
@@ -74,234 +70,6 @@ def process_file(uploaded_file):
         st.error(f"📄 Error processing file: {str(e)}")
         return ""
 
-# Custom model definition
-# Masking functions
-def subsequent_mask(size):
-    attn_shape = (1, size, size)
-    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
-    return torch.from_numpy(subsequent_mask) == 0
-
-def make_std_mask(tgt, pad):
-    tgt_mask = (tgt != pad).unsqueeze(-2)
-    return tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)
-
-# Batch class
-class Batch:
-    def __init__(self, src, trg=None, pad=0):
-        src = torch.from_numpy(src).to(DEVICE).long()
-        self.src = src
-        self.src_mask = (src != pad).unsqueeze(-2)
-        if trg is not None:
-            trg = torch.from_numpy(trg).to(DEVICE).long()
-            self.trg = trg[:, :-1]
-            self.trg_y = trg[:, 1:]
-            self.trg_mask = make_std_mask(self.trg, pad)
-            self.ntokens = (self.trg_y != pad).data.sum()
-
-# Hugging Face config
-class En2FrConfig(PretrainedConfig):
-    model_type = "en2fr_transformer"
-    def __init__(self, src_vocab=32000, tgt_vocab=32000, N=6, d_model=512, 
-                 d_ff=2048, h=8, dropout=0.1, **kwargs):
-        self.src_vocab = src_vocab
-        self.tgt_vocab = tgt_vocab
-        self.N = N
-        self.d_model = d_model
-        self.d_ff = d_ff
-        self.h = h
-        self.dropout = dropout
-        super().__init__(**kwargs)
-
-# Transformer components
-class Transformer(nn.Module):
-    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
-        super().__init__()
-        self.encoder = encoder
-        self.decoder = decoder
-        self.src_embed = src_embed
-        self.tgt_embed = tgt_embed
-        self.generator = generator
-
-    def forward(self, src, tgt, src_mask, tgt_mask):
-        memory = self.encoder(self.src_embed(src), src_mask)
-        output = self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
-        return output
-
-class Encoder(nn.Module):
-    def __init__(self, layer, N):
-        super().__init__()
-        self.layers = nn.ModuleList([deepcopy(layer) for _ in range(N)])
-        self.norm = LayerNorm(layer.size)
-
-    def forward(self, x, mask):
-        for layer in self.layers:
-            x = layer(x, mask)
-        return self.norm(x)
-
-class EncoderLayer(nn.Module):
-    def __init__(self, size, self_attn, feed_forward, dropout):
-        super().__init__()
-        self.self_attn = self_attn
-        self.feed_forward = feed_forward
-        self.sublayer = nn.ModuleList([deepcopy(SublayerConnection(size, dropout)) for _ in range(2)])
-        self.size = size
-
-    def forward(self, x, mask):
-        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
-        return self.sublayer[1](x, self.feed_forward)
-
-class Decoder(nn.Module):
-    def __init__(self, layer, N):
-        super().__init__()
-        self.layers = nn.ModuleList([deepcopy(layer) for _ in range(N)])
-        self.norm = LayerNorm(layer.size)
-
-    def forward(self, x, memory, src_mask, tgt_mask):
-        for layer in self.layers:
-            x = layer(x, memory, src_mask, tgt_mask)
-        return self.norm(x)
-
-class DecoderLayer(nn.Module):
-    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
-        super().__init__()
-        self.size = size
-        self.self_attn = self_attn
-        self.src_attn = src_attn
-        self.feed_forward = feed_forward
-        self.sublayer = nn.ModuleList([deepcopy(SublayerConnection(size, dropout)) for _ in range(3)])
-
-    def forward(self, x, memory, src_mask, tgt_mask):
-        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
-        x = self.sublayer[1](x, lambda x: self.src_attn(x, memory, memory, src_mask))
-        return self.sublayer[2](x, self.feed_forward)
-
-class SublayerConnection(nn.Module):
-    def __init__(self, size, dropout):
-        super().__init__()
-        self.norm = LayerNorm(size)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x, sublayer):
-        return x + self.dropout(sublayer(self.norm(x)))
-
-class LayerNorm(nn.Module):
-    def __init__(self, features, eps=1e-6):
-        super().__init__()
-        self.a_2 = nn.Parameter(torch.ones(features))
-        self.b_2 = nn.Parameter(torch.zeros(features))
-        self.eps = eps
-
-    def forward(self, x):
-        mean = x.mean(-1, keepdim=True)
-        std = x.std(-1, keepdim=True)
-        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
-
-class MultiHeadedAttention(nn.Module):
-    def __init__(self, h, d_model, dropout=0.1):
-        super().__init__()
-        assert d_model % h == 0
-        self.d_k = d_model // h
-        self.h = h
-        self.linears = nn.ModuleList([deepcopy(nn.Linear(d_model, d_model)) for _ in range(4)])
-        self.attn = None
-        self.dropout = nn.Dropout(p=dropout)
-
-    def forward(self, query, key, value, mask=None):
-        if mask is not None:
-            mask = mask.unsqueeze(1)
-        nbatches = query.size(0)
-        query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
-                             for l, x in zip(self.linears, (query, key, value))]
-        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
-        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
-        return self.linears[-1](x)
-
-def attention(query, key, value, mask=None, dropout=None):
-    d_k = query.size(-1)
-    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
-    if mask is not None:
-        scores = scores.masked_fill(mask == 0, -1e9)
-    p_attn = nn.functional.softmax(scores, dim=-1)
-    if dropout is not None:
-        p_attn = dropout(p_attn)
-    return torch.matmul(p_attn, value), p_attn
-
-class PositionwiseFeedForward(nn.Module):
-    def __init__(self, d_model, d_ff, dropout=0.1):
-        super().__init__()
-        self.w_1 = nn.Linear(d_model, d_ff)
-        self.w_2 = nn.Linear(d_ff, d_model)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        return self.w_2(self.dropout(self.w_1(x)))
-
-class Embeddings(nn.Module):
-    def __init__(self, d_model, vocab):
-        super().__init__()
-        self.lut = nn.Embedding(vocab, d_model)
-        self.d_model = d_model
-
-    def forward(self, x):
-        return self.lut(x) * math.sqrt(self.d_model)
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, dropout, max_len=5000):
-        super().__init__()
-        self.dropout = nn.Dropout(p=dropout)
-        pe = torch.zeros(max_len, d_model, device=DEVICE)
-        position = torch.arange(0., max_len, device=DEVICE).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0., d_model, 2, device=DEVICE) * -(math.log(10000.0) / d_model))
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)
-        self.register_buffer('pe', pe)
-
-    def forward(self, x):
-        x = x + self.pe[:, :x.size(1)].requires_grad_(False)
-        return self.dropout(x)
-
-class Generator(nn.Module):
-    def __init__(self, d_model, vocab):
-        super().__init__()
-        self.proj = nn.Linear(d_model, vocab)
-
-    def forward(self, x):
-        return nn.functional.log_softmax(self.proj(x), dim=-1)
-
-def create_model(src_vocab, tgt_vocab, N, d_model, d_ff, h, dropout=0.1):
-    attn = MultiHeadedAttention(h, d_model).to(DEVICE)
-    ff = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)
-    pos = PositionalEncoding(d_model, dropout).to(DEVICE)
-    model = Transformer(
-        Encoder(EncoderLayer(d_model, deepcopy(attn), deepcopy(ff), dropout).to(DEVICE), N).to(DEVICE),
-        Decoder(DecoderLayer(d_model, deepcopy(attn), deepcopy(attn), deepcopy(ff), dropout).to(DEVICE), N).to(DEVICE),
-        nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE), deepcopy(pos)),
-        nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE), deepcopy(pos)),
-        Generator(d_model, tgt_vocab)).to(DEVICE)
-    for p in model.parameters():
-        if p.dim() > 1:
-            nn.init.xavier_uniform_(p)
-    return model
-
-class En2FrTransformer(PreTrainedModel):
-    config_class = En2FrConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = create_model(
-            src_vocab=config.src_vocab,
-            tgt_vocab=config.tgt_vocab,
-            N=config.N,
-            d_model=config.d_model,
-            d_ff=config.d_ff,
-            h=config.h,
-            dropout=config.dropout
-        )
-
-    def forward(self, src, tgt, src_mask, tgt_mask):
-        return self.model(src, tgt, src_mask, tgt_mask)
-
 # Model loading function
 @st.cache_resource
 def load_model(hf_token):
@@ -312,15 +80,18 @@ def load_model(hf_token):
         
         login(token=hf_token)
         
-        # Load tokenizer (assuming a tokenizer was saved with the model)
+        # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
         
-        # Load the custom model
-        model = En2FrTransformer.from_pretrained(
+        # Load the model with appropriate dtype for CPU/GPU compatibility
+        dtype = torch.float16 if DEVICE == "cuda" else torch.float32
+        model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            token=hf_token
+            token=hf_token,
+            torch_dtype=dtype,
+            device_map="auto",  # Automatically maps to CPU or GPU
+            quantization_config=None  # Disable FP8 quantization
         )
-        model.to(DEVICE)  # Ensure model is on the correct device
         
         return model, tokenizer
         
@@ -328,40 +99,39 @@ def load_model(hf_token):
         st.error(f"🤖 Model loading failed: {str(e)}")
         return None
 
-# Simple tokenization function (placeholder, since we don't have the actual vocab)
-def tokenize_text(text, tokenizer, max_length=10):
-    # This is a placeholder; in a real scenario, you'd use the tokenizer's vocabulary
-    # For now, we'll create dummy token IDs (0 for padding, 1 for start, 2 for end, 3+ for words)
-    words = text.split()
-    token_ids = [1] + [i + 3 for i in range(min(len(words), max_length - 2))] + [2]
-    if len(token_ids) < max_length:
-        token_ids += [0] * (max_length - len(token_ids))
-    return torch.tensor([token_ids], dtype=torch.long, device=DEVICE)
-
-# Generation function for translation (custom inference loop)
+# Generation function for translation
 def generate_translation(input_text, model, tokenizer):
-    model.eval()
-    with torch.no_grad():
-        # Tokenize input (source) and target (start with a dummy start token)
-        src = tokenize_text(input_text, tokenizer)
-        tgt = torch.tensor([[1]], dtype=torch.long, device=DEVICE)  # Start token
-        src_mask = (src != 0).unsqueeze(-2)
-        max_length = 10  # Adjust as needed
+    try:
+        # Prepare the prompt
+        full_prompt = TRANSLATION_PROMPT.format(input_text=input_text)
+        
+        # Tokenize the input
+        inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        inputs = inputs.to(DEVICE)
+        
+        # Generate translation
+        model.eval()
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                max_new_tokens=512,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                do_sample=True,
+                num_return_sequences=1
+            )
         
-        # Generate translation token by token
-        for _ in range(max_length - 1):
-            tgt_mask = make_std_mask(tgt, pad=0)
-            output = model(src, tgt, src_mask, tgt_mask)
-            output = model.model.generator(output[:, -1, :])  # Get logits for the last token
-            next_token = torch.argmax(output, dim=-1).unsqueeze(0)
-            tgt = torch.cat((tgt, next_token), dim=1)
-            if next_token.item() == 2:  # End token
-                break
+        # Decode the output
+        translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Convert token IDs back to text (placeholder)
-        # In a real scenario, you'd use tokenizer.decode()
-        translation = " ".join([f"word{i-3}" if i >= 3 else "<start>" if i == 1 else "<end>" for i in tgt[0].tolist()])
+        # Extract the French translation part (after the prompt)
+        translation = translation.split("**French translation:**")[-1].strip()
         return translation
+        
+    except Exception as e:
+        raise Exception(f"Generation error: {str(e)}")
 
 # Display chat messages
 for message in st.session_state.messages:
@@ -411,10 +181,10 @@ if prompt := st.chat_input("Enter text to translate into French..."):
                 st.markdown(translation)
                 st.session_state.messages.append({"role": "assistant", "content": translation})
                 
-                # Calculate performance metrics (simplified, since we don't have real token counts)
+                # Calculate performance metrics
                 end_time = time.time()
-                input_tokens = len(input_text.split())  # Approximate
-                output_tokens = len(translation.split())  # Approximate
+                input_tokens = len(tokenizer(input_text)["input_ids"])
+                output_tokens = len(tokenizer(translation)["input_ids"])
                 speed = output_tokens / (end_time - start_time)
                 
                 # Calculate costs (hypothetical pricing model)