from django.shortcuts import render
from django.http import JsonResponse
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# טעינת המודל והטוקניזר של DistilGPT-2
model_name = "distilgpt2"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# הגדרת pad_token כדי למנוע בעיות עם ריווח
tokenizer.pad_token = tokenizer.eos_token

def chat_view(request):
    return render(request, 'chat/chat.html')

def generate_response(request):
    if request.method == "POST":
        user_input = request.POST.get("prompt", "")

        # טוקניזציה של הקלט
        inputs = tokenizer(user_input, return_tensors='pt', padding=True, truncation=True)
        attention_mask = inputs['attention_mask']
        input_ids = inputs['input_ids']

        # יצירת טקסט באמצעות המודל
        outputs = model.generate(input_ids, attention_mask=attention_mask, max_length=200, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)

        # פענוח התשובה שנוצרה
        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)

        return JsonResponse({"response": response_text})

    return JsonResponse({"response": "Invalid request method."})

from django.urls import path
from .views import chat_view, generate_response

urlpatterns = [
    path('', chat_view, name='chat'),
    path('generate_response/', generate_response, name='generate_response'),
]