rollback

app.py

@@ -1,5 +1,4 @@
 import os
-import torch
 from flask import Flask, jsonify, request
 from flask_cors import CORS 
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig  
@@ -15,7 +14,6 @@ CORS(app, resources={r"api/predict/*": {"origins": ["http://localhost:3000", "ht
 # Global variables for model and tokenizer
 model = None
 tokenizer = None
-device = torch.device("cpu")  # Explicitly use CPU
 
 def get_model_and_tokenizer(model_id):
     global model, tokenizer
@@ -25,12 +23,11 @@ def get_model_and_tokenizer(model_id):
         tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
         tokenizer.pad_token = tokenizer.eos_token
         
-        print(f"Loading model for model_id: {model_id}")
-        # Load the model and move it to the specified device
-        model = AutoModelForCausalLM.from_pretrained(model_id).to(device) #, device_map="auto")
+        print(f"Loading model and for model_id: {model_id}")
+        # Load the model
+        model = AutoModelForCausalLM.from_pretrained(model_id) #, device_map="auto")
         model.config.use_cache = False
-        print("Model and tokenizer successfully loaded.")
-
+   
     except Exception as e:
         print(f"Error loading model: {e}")
 
@@ -41,13 +38,13 @@ def generate_response(user_input, model_id):
     if model is None or tokenizer is None:
         get_model_and_tokenizer(model_id)  # Load model and tokenizer
 
-    # Prepare the input tensors and move them to the appropriate device
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    # Prepare the input tensors
+    inputs = tokenizer(prompt, return_tensors="pt")  # Move inputs to GPU if available
     
     generation_config = GenerationConfig(
         max_new_tokens=100,
         min_length=5,
-        #temperature=0.7, 
+        temperature=0.7, 
         do_sample=False,
         num_beams=1,
         pad_token_id=tokenizer.eos_token_id,