model run for mode_response

app.py

@@ -5,7 +5,7 @@ os.environ["HF_HOME"] = "/workspace/huggingface_cache"  # Change this to a writa
 
 from flask import Flask, jsonify, request
 from flask_cors import CORS 
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig  
 
 app = Flask(__name__)
 
@@ -13,16 +13,43 @@ app = Flask(__name__)
 CORS(app, resources={r"api/predict/*": {"origins": ["http://localhost:3000", "https://main.dbn2ikif9ou3g.amplifyapp.com"]}})
 
 # Model setup
-model_id = "YALCINKAYA/opsgenius-large"
+model_id = "YALCINKAYA/opsgenius-large" 
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+def get_model_and_tokenizer(model_id):
+    # Load the tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load the model
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+    model.config.use_cache = False
+    return model, tokenizer
+
+model, tokenizer = get_model_and_tokenizer(model_id)
 
 def generate_response(user_input):
+
+    prompt = formatted_prompt(user_input)
+    # Prepare the input tensors
+    inputs = tokenizer(prompt, return_tensors="pt")#.to('cuda')  # Move inputs to GPU
+   
+    generation_config = GenerationConfig(
+        max_new_tokens=100,  # Allow enough length for full responses
+        min_length=5,
+        temperature=0.7, 
+        do_sample=False,  # Set to False for deterministic responses 
+        num_beams=1,
+        pad_token_id=tokenizer.eos_token_id,  # Set pad_token_id
+        truncation=True  # Enable truncation
+    )
+
     # Instead of generating a response from the model, return a dummy message
-    dummy_response = "This is a dummy response for the input: " + user_input
-    return dummy_response
+    #dummy_response = "This is a dummy response for the input: " + user_input
+    # Generate response
+    outputs = model.generate(**inputs, generation_config=generation_config)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return response
 
 def formatted_prompt(question) -> str:
     return f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant:"