Update app.py

app.py

@@ -2,7 +2,7 @@ from flask import Flask, render_template, request, jsonify
 import subprocess
 import tempfile
 import os
-from langchain.llms import Ollama
+from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain.memory import ConversationBufferMemory
@@ -13,6 +13,8 @@ import sqlite3
 from contextlib import contextmanager
 import re
 from werkzeug.utils import secure_filename
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 app = Flask(__name__)
 
@@ -32,27 +34,27 @@ import torch
 # Load model and tokenizer
 model_name = "mistralai/Mistral-7B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-class HuggingFaceLLM:
-    def __init__(self, model, tokenizer):
-        self.model = model
-        self.tokenizer = tokenizer
-    
-    def predict(self, prompt):
-        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=2048, truncation=True)
-        with torch.no_grad():
-            outputs = self.model.generate(
-                inputs["input_ids"],
-                max_length=2048,
-                num_return_sequences=1,
-                temperature=0.7,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    load_in_8bit=True
+)
+
+# Create pipeline
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=512,
+    temperature=0.7,
+    top_p=0.95,
+    repetition_penalty=1.15
+)
+
+
 
-llm = HuggingFaceLLM(model, tokenizer)
+llm = HuggingFacePipeline(pipeline=pipe)
 
 @contextmanager
 def get_db_connection():