Update src/models/flan.py

src/models/flan.py

@@ -1,4 +1,3 @@
-# models/flan.py
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
@@ -9,7 +8,7 @@ class FlanT5:
         self.model = AutoModelForSeq2SeqLM.from_pretrained(
             model_id,
             device_map="auto"
-        )
+        ).eval()
 
     def generate(self, prompt, temperature=0.7):
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
@@ -18,6 +17,8 @@ class FlanT5:
             do_sample=True,
             temperature=temperature,
             top_p=0.9,
-            max_new_tokens=100
+            max_new_tokens=64,
+            no_repeat_ngram_size=2,
+            eos_token_id=self.tokenizer.eos_token_id
         )
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)