Update app.py

app.py

@@ -7,6 +7,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from deep_translator import GoogleTranslator
 import torch
 
+# تنظیم متغیر محیطی برای دیباگ CUDA
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+
 # کلید COMTRADE
 subscription_key = os.getenv("COMTRADE_API_KEY", "")
 # توکن Hugging Face
@@ -19,14 +22,22 @@ translator = GoogleTranslator(source='en', target='fa')
 quantization_config = BitsAndBytesConfig(load_in_4bit=True)
 
 # بارگذاری توکنایزر و مدل
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-27b-it", token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-3-27b-it",
-    token=hf_token,
-    quantization_config=quantization_config,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
+try:
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", token=hf_token)
+    model = AutoModelForCausalLM.from_pretrained(
+        "google/gemma-2b-it",
+        token=hf_token,
+        quantization_config=quantization_config,
+        device_map="auto",
+        torch_dtype=torch.float16
+    )
+except Exception as e:
+    print(f"خطا در بارگذاری مدل: {str(e)}")
+    raise e
+
+# تنظیم صریح pad_token_id
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
 
 # تابع دریافت اطلاعات واردکنندگان
 def get_importers(hs_code: str, year: str, month: str):
@@ -50,32 +61,37 @@ def get_importers(hs_code: str, year: str, month: str):
     return result
 
 # تابع ارائه مشاوره با استفاده از GPU
-@spaces.GPU(duration=180)  # افزایش مدت زمان برای مدل سنگین
+@spaces.GPU(duration=120)
 def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: str):
     if table_data is None or table_data.empty:
         return "ابتدا باید اطلاعات واردات را نمایش دهید."
     
     table_str = table_data.to_string(index=False)
     period = f"{year}/{int(month):02d}"
+    # پرامپت بهینه‌شده
     prompt = (
-        f"The following table shows countries that imported a product with HS code {hs_code} during the period {period}:\n"
-        f"{table_str}\n\n"
-        f"Please provide a detailed and comprehensive analysis in two paragraphs. The first paragraph should discuss market opportunities, potential demand, and specific cultural or economic factors influencing the demand for this product in these countries. The second paragraph should offer actionable strategic recommendations for exporters, including detailed trade strategies, risk management techniques, and steps to establish local partnerships."
+        f"Table of countries importing HS code {hs_code} in {period}:\n{table_str}\n\n"
+        f"Analyze market opportunities and cultural/economic factors in one paragraph. "
+        f"Provide strategic recommendations for exporters in another paragraph."
     )
     print("پرامپت ساخته‌شده:")
     print(prompt)
 
     try:
         # آماده‌سازی ورودی برای مدل
-        input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        input_ids = inputs.input_ids.to("cuda")
+        attention_mask = inputs.attention_mask.to("cuda")
+        
         # تولید خروجی
         outputs = model.generate(
-            **input_ids,
-            max_new_tokens=1024,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=512,
             do_sample=True,
-            temperature=0.6,  # برای پاسخ‌های منسجم
-            top_p=0.85,      # برای کیفیت بهتر
-            pad_token_id=tokenizer.eos_token_id  # جلوگیری از خطای pad token
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
         )
         # دیکد کردن خروجی و حذف پرامپت
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)