Upload app.py

app.py

@@ -10,6 +10,8 @@ import os
 
 # Load the model from Hugging Face Hub
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
 
 # Define the SmolLM2-135M model (a simplified version of a Transformer)
 class SmolLM(nn.Module):
@@ -64,38 +66,13 @@ model = load_model()
 model.train(False)
 
 def generate_text(prompt, max_length=100, num_samples=1, temperature=0.8):
-    enc = tiktoken.get_encoding('gpt2')
-    tokens = enc.encode(prompt)
-    tokens = torch.tensor(tokens, dtype=torch.long)
-    tokens = tokens.unsqueeze(0).repeat(num_samples, 1)
-    tokens = tokens.to(device)
-    
-    with torch.no_grad():
-        for _ in range(max_length):
-            if tokens.size(1) >= 1024:  # GPT context length
-                break
-                
-            logits = model(tokens)[0]
-            logits = logits[:, -1, :]
-            #logits = logits[:, -1, :] / temperature
-            probs = F.softmax(logits, dim=-1)
-            
-            # Top-k sampling
-            topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
-            ix = torch.multinomial(topk_probs, 1)
-            next_token = torch.gather(topk_indices, -1, ix)
-            
-            tokens = torch.cat((tokens, next_token), dim=1)
-            
-            # Remove special token check entirely
-            # Just generate for the specified length or until context limit
-    
-    generated_texts = []
-    for i in range(num_samples):
-        text = enc.decode(tokens[i].tolist())
-        generated_texts.append(text)
-    
-    return '\n\n---\n\n'.join(generated_texts)
+
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    outputs = model(input_ids)
+    predictions = torch.argmax(outputs, dim=-1)
+    decoded = tokenizer.decode(predictions[0], skip_special_tokens=True)
+    return decoded
+  
 
 # Create Gradio interface
 iface = gr.Interface(