Update app.py

app.py

@@ -1,7 +1,8 @@
 from fastapi import FastAPI, Request
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from fastapi.responses import StreamingResponse
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import torch.nn.functional as F
 
 app = FastAPI()
 
@@ -9,6 +10,8 @@ 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)
 
 @app.get("/")
 def read_root():
@@ -22,22 +25,49 @@ async def predict(request: Request):
         return {"error": "Prompt is required"}
 
     # Tokenize the input
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")  # Use "cuda" if GPU is enabled
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    input_ids = inputs.input_ids
+    attention_mask = inputs.attention_mask
 
     # Generator function to stream tokens
     def token_generator():
-        outputs = model.generate(
-            inputs.input_ids,
-            max_length=100,
-            do_sample=True,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9,
-            repetition_penalty=1.2,
-        )
-        for token_id in outputs[0]:
-            token = tokenizer.decode(token_id, skip_special_tokens=True)
-            yield f"{token} "
-    
-    # Return StreamingResponse
+        temperature = 0.7
+        top_p = 0.9
+
+        for _ in range(100):  # Limit the number of generated tokens
+            # Get the model outputs
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            next_token_logits = outputs.logits[:, -1, :]  # Logits for the last token
+
+            # Apply temperature scaling
+            next_token_logits = next_token_logits / temperature
+
+            # Convert logits to probabilities
+            next_token_probs = F.softmax(next_token_logits, dim=-1)
+
+            # Apply top-p nucleus 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 from the filtered distribution
+            if len(sorted_probs) > 0:
+                next_token_id = sorted_indices[torch.multinomial(sorted_probs, 1)]
+            else:
+                # Fallback to greedy selection if no tokens meet top-p
+                next_token_id = torch.argmax(next_token_probs)
+
+            # Append the generated token to the input
+            input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
+
+            # Decode the token and yield it
+            token = tokenizer.decode(next_token_id.squeeze(), skip_special_tokens=True)
+            yield token + " "
+
+            # Stop if the model generates the end-of-sequence token
+            if next_token_id.squeeze().item() == tokenizer.eos_token_id:
+                break
+
+    # Return the generator as a streaming response
     return StreamingResponse(token_generator(), media_type="text/plain")