Update main.py

main.py

@@ -12,7 +12,6 @@ test_model = Llama(model_name)
 
 class RequestBody(BaseModel):
     prompt: str
-    num_return_sequences: int = 10
     temperature: float = 1.0
     top_k: int = 50
     top_p: float = 1.0
@@ -21,13 +20,12 @@ class RequestBody(BaseModel):
 async def generate_text(request: RequestBody):
     try:
         prompt = sf.encoder(request.prompt)
-        input_ids = test_tokenizer(prompt, return_tensors='pt', truncation=False).input_ids
-        outputs = test_model.generate(
-            input_ids=input_ids, 
+        outputs = test_model(
+            prompt,
             max_new_tokens=512, 
             num_beams=10, 
             early_stopping=True, 
-            num_return_sequences=request.num_return_sequences, 
+            num_return_sequences=10, 
             do_sample=True,
             top_k = request.top_k,
             top_p = request.top_p,
@@ -35,8 +33,8 @@ async def generate_text(request: RequestBody):
         )
         
         result = {'input': prompt}
-        for i in range(num_return_sequences):
-            output1 = test_tokenizer.batch_decode(outputs.detach().numpy(), skip_special_tokens=True)[i][len(prompt):]
+        for i in range(10):
+            output1 = outputs[i][len(prompt):]
             first_inst_index = output1.find("[/INST]")
             second_inst_index = output1.find("[/IN", first_inst_index + len("[/INST]") + 1)
             predicted_selfies = output1[first_inst_index + len("[/INST]"):second_inst_index].strip()