Update app.py

app.py

@@ -10,6 +10,7 @@ app = FastAPI()
 model_name = "EleutherAI/gpt-neo-1.3B"  # Replace with your desired model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
@@ -20,45 +21,60 @@ async def predict(request: Request):
     if not prompt:
         return {"error": "Prompt is required"}
 
-    # Tokenize the input
+    # Tokenize the input and move to correct device
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     input_ids = inputs.input_ids
     attention_mask = inputs.attention_mask
 
     def token_generator():
+        # Use nonlocal to allow reassigning input_ids inside the nested function
+        nonlocal input_ids
+
+        # Sampling parameters
         temperature = 0.7
         top_p = 0.9
+        max_new_tokens = 30
 
-        for _ in range(100):  # Limit to 100 tokens
-            with torch.no_grad():  # Disable gradient computation for inference
+        for _ in range(max_new_tokens):
+            with torch.no_grad():
+                # Forward pass
                 outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                 next_token_logits = outputs.logits[:, -1, :]
 
-                # Apply temperature and softmax
+                # Temperature scaling
                 next_token_logits = next_token_logits / temperature
+
+                # Convert logits to probabilities
                 next_token_probs = F.softmax(next_token_logits, dim=-1)
 
-                # Apply nucleus sampling (top-p)
+                # Apply nucleus (top-p) sampling
                 sorted_probs, sorted_indices = torch.sort(next_token_probs, descending=True)
                 cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
-                sorted_probs = sorted_probs[cumulative_probs <= top_p]
-                sorted_indices = sorted_indices[:len(sorted_probs)]
 
-                # Sample next token
-                if len(sorted_probs) > 0:
-                    next_token_id = sorted_indices[torch.multinomial(sorted_probs, 1)]
+                # Filter out tokens above the top_p threshold
+                valid_indices = cumulative_probs <= top_p
+                filtered_probs = sorted_probs[valid_indices]
+                filtered_indices = sorted_indices[valid_indices]
+
+                if len(filtered_probs) == 0:
+                    # Fallback to greedy if no tokens meet top_p
+                    next_token_id = torch.argmax(next_token_probs).unsqueeze(-1)
                 else:
-                    next_token_id = torch.argmax(next_token_probs)
+                    # Sample from the filtered distribution
+                    sampled_id = torch.multinomial(filtered_probs, num_samples=1)
+                    next_token_id = filtered_indices[sampled_id].unsqueeze(-1)
 
-                # Append the new token to the input sequence
-                input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
+                # Append the new token to our running sequence
+                input_ids = torch.cat([input_ids, next_token_id], dim=-1)
 
                 # Decode and yield the token
                 token = tokenizer.decode(next_token_id.squeeze(), skip_special_tokens=True)
                 yield token + " "
 
-                # Stop if the end-of-sequence token is generated
-                if next_token_id.squeeze().item() == tokenizer.eos_token_id:
-                    break
+                # Stop if EOS token is generated
+                if tokenizer.eos_token_id is not None:
+                    if next_token_id.squeeze().item() == tokenizer.eos_token_id:
+                        break
 
+    # Return the streaming response
     return StreamingResponse(token_generator(), media_type="text/plain")