Fixed app v2

app.py

@@ -3,20 +3,20 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
-# Define model paths
-model_path = "/home/migel/Llama-3.1-WhiteRabbitNeo-2-8B"
+# Hugging Face model repository path
+model_name = "WhiteRabbitNeo/WhiteRabbitNeo-13B-v1"
 output_file_path = "/home/user/conversations.jsonl"
 
-# Load model and tokenizer
+# Load model and tokenizer from Hugging Face model hub
 model = AutoModelForCausalLM.from_pretrained(
-    model_path,
+    model_name,
     torch_dtype=torch.float16,
     device_map="auto",
     load_in_4bit=False,
     trust_remote_code=False,
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -26,8 +26,7 @@ class PromptRequest(BaseModel):
     user_input: str
 
 def generate_text(instruction):
-    tokens = tokenizer.encode(instruction)
-    tokens = torch.LongTensor(tokens).unsqueeze(0).to("cuda")
+    tokens = tokenizer.encode(instruction, return_tensors="pt").to("cuda")
 
     instance = {
         "input_ids": tokens,
@@ -37,7 +36,7 @@ def generate_text(instruction):
         "top_k": 50,
     }
 
-    length = len(tokens[0])
+    length = tokens.size(1)
     with torch.no_grad():
         rest = model.generate(
             input_ids=tokens,