Update main.py

main.py

@@ -1,6 +1,11 @@
+from transformers import AutoModelForCausalLM, GemmaTokenizerFast, TextIteratorStreamer
 from flask import Flask, request, jsonify
-from hugchat import hugchat
-from hugchat.login import Login
+
+from threading import Thread
+from typing import Iterator
+
+import spaces
+import torch
 import os
 
 app = Flask(__name__)
@@ -10,10 +15,24 @@ print("Hello welcome to Sema AI", flush=True)  # Flush to ensure immediate outpu
 # Get Hugging Face credentials from environment variables
 email = os.getenv('HF_EMAIL')
 password = os.getenv('HF_PASS')
+GEMMA_TOKEN = os.getenv("GEMMA_TOKEN")
+#print(f"email is {email} and password is {password}", flush=True)
+
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-#email = "[email protected]"
-#password = "@Tri@n@th@1"
-print(f"email is {email} and password is {password}", flush=True)
+model_id = "google/gemma-2-2b-it"
+tokenizer = GemmaTokenizerFast.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+model.config.sliding_window = 4096
+model.eval()
 
 @app.route("/")
 def hello():