app.py

from fastapi import FastAPI, Query
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os

app = FastAPI()

# بارگیری مدل و توکنایزر
model_name = "moonshotai/Kimi-VL-A3B-Thinking"
hf_token = os.getenv("HUGGINGFACE_TOKEN")  # خواندن توکن از متغیر محیطی

if hf_token is None:
    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set.")

# بررسی بارگیری مدل
try:
    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    raise

# بررسی توکنایزر
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
    test_text = "This is a test."
    tokens = tokenizer(test_text, return_tensors="pt")
    print("Tokenizer works successfully.")
    print(tokens)
except Exception as e:
    print(f"Error with tokenizer: {e}")
    raise

@app.get("/generate")
async def generate_text(prompt: str = Query(..., description="متن ورودی برای تولید متن")):
    try:
        # توکن‌بندی ورودی
        inputs = tokenizer(prompt, return_tensors="pt")
        # تولید متن با استفاده از مدل
        outputs = model.generate(**inputs, max_length=100)
        # تبدیل خروجی به متن
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return {"generated_text": generated_text}
    except Exception as e:
        return {"error": str(e)}