Update main.py

main.py

@@ -22,7 +22,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 class req(BaseModel):
   prompt: str
-  history: list
 
 @app.get("/")
 def read_root():
@@ -31,18 +30,15 @@ def read_root():
 @app.post("/api")
 def read_root(data: req):
   print("Prompt:", data.prompt)
-  print("History:", data.history)
 
-  history_string = "\n".join(data.history)
-  
   input_text = data.prompt
   
   # Tokenize the input text
-  inputs = tokenizer.encode_plus(history_string, input_text, return_tensors="pt")
+  input_ids = tokenizer.encode(input_text, return_tensors="pt")
   
   # Generate output using the model
-  outputs = model.generate(**inputs)
-  generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+  output_ids = model.generate(input_ids, num_beams=5, no_repeat_ngram_size=2)
+  generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
   
   answer_data = { "answer": generated_text }
   print("Answer:", generated_text)