resolve a issue

generate.py

@@ -64,17 +64,17 @@ def generate_code(model_data, prompt, max_new_tokens):
     input_ids = tokenizer.encode(prompt, return_tensors="pt")
 
     # ایجاد attention mask
-    attention_mask = torch.ones(input_ids.shape, device=input_ids.device)  # ایجاد یک ماسک توجه برای ورودی‌ها
+    attention_mask = torch.ones(input_ids.shape, device=input_ids.device)
 
     # تولید کد
     outputs = model.generate(
         input_ids=input_ids,
-        attention_mask=attention_mask,  # ارسال attention mask
+        attention_mask=attention_mask,
         max_new_tokens=max_new_tokens,
-        do_sample=False,
-        pad_token_id=tokenizer.eos_token_id,  # تنظیم شناسه توکن پایان به عنوان پرکننده
-        repetition_penalty=1.2,  # جلوگیری از تکرار
-        no_repeat_ngram_size=3,  # جلوگیری از تکرار n-gram
+        do_sample=False,  # غیرفعال کردن نمونه‌گیری (حالت حریصانه)
+        pad_token_id=tokenizer.eos_token_id,
+        repetition_penalty=1.2,
+        no_repeat_ngram_size=3,
     )
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)

model.py

@@ -12,7 +12,7 @@ model_dict = {
     "dialoGPT": {"path": "microsoft/DialoGPT-small", "library": AutoModelForCausalLM, "tokenizer": AutoTokenizer, "use_pipeline": False},
     "dialoGPT-medium": {"path": "microsoft/DialoGPT-medium", "library": AutoModelForCausalLM, "tokenizer": AutoTokenizer, "use_pipeline": False},
     "dialoGPT-large": {"path": "microsoft/DialoGPT-large", "library": AutoModelForCausalLM, "tokenizer": AutoTokenizer, "use_pipeline": False},
-    "GPT-Neo-125M": {"path": "EleutherAI/gpt-neo-125m", "library": AutoModelForCausalLM, "tokenizer": AutoTokenizer, "use_pipeline": True},  # اضافه کردن مدل جدید
+    "GPT-Neo-125M": {"path": "EleutherAI/gpt-neo-125m", "library": AutoModelForCausalLM, "tokenizer": AutoTokenizer, "use_pipeline": True},
     "bert-emotion": {"path": "bhadresh-savani/distilbert-base-uncased-emotion", "library": AutoModelForSequenceClassification, "tokenizer": AutoTokenizer, "use_pipeline": True},
 }
 
@@ -31,19 +31,18 @@ def load_model_lazy(model_name):
     if model_info.get("use_pipeline", False):
         print(f"Using pipeline for model: {model_name}")
         if model_name == "bert-emotion":
-            # برای مدل bert-emotion از text-classification استفاده کنید
             model_pipeline = pipeline(
-                "text-classification",  # تغییر وظیفه به text-classification
+                "text-classification",
                 model=model_info["path"],
                 truncation=True
             )
         else:
-            # برای سایر مدل‌ها از text-generation استفاده کنید
             model_pipeline = pipeline(
                 "text-generation",
                 model=model_info["path"],
                 truncation=True,
-                pad_token_id=50256
+                pad_token_id=50256,
+                do_sample=False  # غیرفعال کردن نمونه‌گیری (حالت حریصانه)
             )
         loaded_models[model_name] = {"pipeline": model_pipeline}
         return {"pipeline": model_pipeline}